buildgrid/_app/bots/buildbox.py

@@ -16,82 +16,100 @@
 import os
 import subprocess
 import tempfile
+import uuid
 
 from google.protobuf import any_pb2
 
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
-from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
-from buildgrid.utils import parse_to_pb2_from_fetch
+from buildgrid._protos.google.bytestream import bytestream_pb2, bytestream_pb2_grpc
+from buildgrid.utils import create_digest, parse_to_pb2_from_fetch
+from buildgrid.utils import read_file, write_file
 
 
 def work_buildbox(context, lease):
+    """Executes a lease for a build action, using buildbox.
+    """
+
+    stub_bytestream = bytestream_pb2_grpc.ByteStreamStub(context.cas_channel)
+    local_cas_directory = context.local_cas
     logger = context.logger
 
-    action_digest_any = lease.payload
     action_digest = remote_execution_pb2.Digest()
-    action_digest_any.Unpack(action_digest)
+    lease.payload.Unpack(action_digest)
+
+    action = parse_to_pb2_from_fetch(remote_execution_pb2.Action(),
+                                     stub_bytestream, action_digest)
 
-    stub = bytestream_pb2_grpc.ByteStreamStub(context.cas_channel)
+    command = parse_to_pb2_from_fetch(remote_execution_pb2.Command(),
+                                      stub_bytestream, action.command_digest)
 
-    action = remote_execution_pb2.Action()
-    parse_to_pb2_from_fetch(action, stub, action_digest)
+    environment = dict()
+    for variable in command.environment_variables:
+        if variable.name not in ['PWD']:
+            environment[variable.name] = variable.value
 
-    casdir = context.local_cas
-    remote_command = remote_execution_pb2.Command()
-    parse_to_pb2_from_fetch(remote_command, stub, action.command_digest)
+    if command.working_directory:
+        working_directory = command.working_directory
+    else:
+        working_directory = '/'
 
-    environment = dict((x.name, x.value) for x in remote_command.environment_variables)
     logger.debug("command hash: {}".format(action.command_digest.hash))
     logger.debug("vdir hash: {}".format(action.input_root_digest.hash))
-    logger.debug("\n{}".format(' '.join(remote_command.arguments)))
+    logger.debug("\n{}".format(' '.join(command.arguments)))
+
+    os.makedirs(os.path.join(local_cas_directory, 'tmp'), exist_ok=True)
+    os.makedirs(context.fuse_dir, exist_ok=True)
 
-    # Input hash must be written to disk for buildbox.
-    os.makedirs(os.path.join(casdir, 'tmp'), exist_ok=True)
-    with tempfile.NamedTemporaryFile(dir=os.path.join(casdir, 'tmp')) as input_digest_file:
-        with open(input_digest_file.name, 'wb') as f:
-            f.write(action.input_root_digest.SerializeToString())
-            f.flush()
+    with tempfile.NamedTemporaryFile(dir=os.path.join(local_cas_directory, 'tmp')) as input_digest_file:
+        # Input hash must be written to disk for BuildBox
+        write_file(input_digest_file.name, action.input_root_digest.SerializeToString())
 
-        with tempfile.NamedTemporaryFile(dir=os.path.join(casdir, 'tmp')) as output_digest_file:
-            command = ['buildbox',
+        with tempfile.NamedTemporaryFile(dir=os.path.join(local_cas_directory, 'tmp')) as output_digest_file:
+            command_line = ['buildbox',
                             '--remote={}'.format(context.remote_cas_url),
                             '--input-digest={}'.format(input_digest_file.name),
                             '--output-digest={}'.format(output_digest_file.name),
-                       '--local={}'.format(casdir)]
+                            '--chdir={}'.format(working_directory),
+                            '--local={}'.format(local_cas_directory)]
 
             if context.cas_client_key:
-                command.append('--client-key={}'.format(context.cas_client_key))
+                command_line.append('--client-key={}'.format(context.cas_client_key))
             if context.cas_client_cert:
-                command.append('--client-cert={}'.format(context.cas_client_cert))
+                command_line.append('--client-cert={}'.format(context.cas_client_cert))
             if context.cas_server_cert:
-                command.append('--server-cert={}'.format(context.cas_server_cert))
+                command_line.append('--server-cert={}'.format(context.cas_server_cert))
 
-            if 'PWD' in environment and environment['PWD']:
-                command.append('--chdir={}'.format(environment['PWD']))
+            command_line.append(context.fuse_dir)
+            command_line.extend(command.arguments)
 
-            command.append(context.fuse_dir)
-            command.extend(remote_command.arguments)
-
-            logger.debug(' '.join(command))
+            logger.debug(' '.join(command_line))
             logger.debug("Input root digest:\n{}".format(action.input_root_digest))
             logger.info("Launching process")
 
-            proc = subprocess.Popen(command,
+            command_line = subprocess.Popen(command_line,
                                             stdin=subprocess.PIPE,
                                             stdout=subprocess.PIPE)
-            proc.communicate()
+            # TODO: Should return the stdout and stderr to the user.
+            command_line.communicate()
+
+            output_digest = remote_execution_pb2.Digest()
+            output_digest.ParseFromString(read_file(output_digest_file.name))
 
-            output_root_digest = remote_execution_pb2.Digest()
-            with open(output_digest_file.name, 'rb') as f:
-                output_root_digest.ParseFromString(f.read())
-            logger.debug("Output root digest: {}".format(output_root_digest))
+            logger.debug("Output root digest: {}".format(output_digest))
 
-            if len(output_root_digest.hash) < 64:
+            if len(output_digest.hash) < 64:
                 logger.warning("Buildbox command failed - no output root digest present.")
-            output_file = remote_execution_pb2.OutputDirectory(tree_digest=output_root_digest)
+
+            # TODO: Have BuildBox helping us creating the Tree instance here
+            # See https://gitlab.com/BuildStream/buildbox/issues/7 for details
+            output_tree_digest = _cas_tree_maker(stub_bytestream, output_digest)
+
+            output_directory = remote_execution_pb2.OutputDirectory()
+            output_directory.tree_digest.CopyFrom(output_tree_digest)
+            output_directory.path = os.path.relpath(working_directory, start='/')
 
     action_result = remote_execution_pb2.ActionResult()
-    action_result.output_directories.extend([output_file])
+    action_result.output_directories.extend([output_directory])
 
     action_result_any = any_pb2.Any()
     action_result_any.Pack(action_result)
@@ -99,3 +117,40 @@ def work_buildbox(context, lease):
     lease.result.CopyFrom(action_result_any)
 
     return lease
+
+
+def _cas_tree_maker(stub_bytestream, directory_digest):
+    # Generates and stores a Tree for a given Directory
+    output_tree = remote_execution_pb2.Tree()
+
+    def list_directories(parent_directory):
+        directory_list = list()
+        for directory_node in parent_directory.directories:
+            directory = parse_to_pb2_from_fetch(remote_execution_pb2.Directory(),
+                                                stub_bytestream, directory_node.digest)
+            directory_list.extend(list_directories(directory))
+            directory_list.append(directory)
+
+        return directory_list
+
+    root_directory = parse_to_pb2_from_fetch(remote_execution_pb2.Directory(),
+                                             stub_bytestream, directory_digest)
+    output_tree.children.extend(list_directories(root_directory))
+    output_tree.root.CopyFrom(root_directory)
+
+    output_tree_blob = output_tree.SerializeToString()
+    output_tree_digest = create_digest(output_tree_blob)
+
+    # XXX: BuildBox does not support instance name for now...
+    resource_name = '/'.join(['uploads', str(uuid.uuid4()),
+                              'blobs', output_tree_digest.hash, str(output_tree_digest.size_bytes)])
+
+    write_request = bytestream_pb2.WriteRequest(resource_name=resource_name)
+    write_request.data = output_tree_blob
+    write_request.finish_write = True
+
+    write_response = stub_bytestream.Write(iter([write_request]))
+
+    assert write_response.committed_size == len(output_tree_blob)
+
+    return output_tree_digest

buildgrid/utils.py

@@ -243,6 +243,21 @@ def read_file(file_path):
         return byte_file.read()
 
 
+def write_file(file_path, content):
+    """Dumps raw memory content to a file.
+
+    Args:
+        file_path (str): path to the target file.
+        content (bytes): raw file's content.
+
+    Raises:
+        OSError: If `file_path` does not exist or is not writable.
+    """
+    with open(file_path, 'wb') as byte_file:
+        byte_file.write(content)
+        byte_file.flush()
+
+
 def output_file_maker(file_path, input_path):
     """Creates an :obj:`OutputFile` from a local file.