temp_directory.py: Better action execution support

buildgrid/_app/bots/temp_directory.py

@@ -19,7 +19,7 @@ import tempfile
 
 from google.protobuf import any_pb2
 
-from buildgrid.utils import read_file, create_digest, write_fetch_directory, parse_to_pb2_from_fetch
+from buildgrid.utils import output_file_maker, write_fetch_directory, parse_to_pb2_from_fetch
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
 from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
 
@@ -35,54 +35,65 @@ def work_temp_directory(context, lease):
     action_digest = remote_execution_pb2.Digest()
     lease.payload.Unpack(action_digest)
 
-    action = remote_execution_pb2.Action()
+    action = parse_to_pb2_from_fetch(remote_execution_pb2.Action(),
+                                     stub_bytestream, action_digest, instance_name)
 
-    action = parse_to_pb2_from_fetch(action, stub_bytestream, action_digest, instance_name)
+    with tempfile.TemporaryDirectory() as temp_directory:
 
-    with tempfile.TemporaryDirectory() as temp_dir:
+        command = parse_to_pb2_from_fetch(remote_execution_pb2.Command(),
+                                          stub_bytestream, action.command_digest, instance_name)
 
-        command = remote_execution_pb2.Command()
-        command = parse_to_pb2_from_fetch(command, stub_bytestream, action.command_digest, instance_name)
+        write_fetch_directory(temp_directory, stub_bytestream,
+                              action.input_root_digest, instance_name)
 
-        arguments = "cd {} &&".format(temp_dir)
+        execution_envionment = os.environ.copy()
+        for variable in command.environment_variables:
+            if variable.name not in ['PATH', 'PWD']:
+                execution_envionment[variable.name] = variable.value
 
+        command_arguments = list()
         for argument in command.arguments:
-            arguments += " {}".format(argument)
-
-        context.logger.info(arguments)
-
-        write_fetch_directory(temp_dir, stub_bytestream, action.input_root_digest, instance_name)
-
-        proc = subprocess.Popen(arguments,
-                                shell=True,
-                                stdin=subprocess.PIPE,
-                                stdout=subprocess.PIPE)
-
-        # TODO: Should return the std_out to the user
-        proc.communicate()
-
-        result = remote_execution_pb2.ActionResult()
-        requests = []
-        for output_file in command.output_files:
-            path = os.path.join(temp_dir, output_file)
-            chunk = read_file(path)
-
-            digest = create_digest(chunk)
-
-            result.output_files.extend([remote_execution_pb2.OutputFile(path=output_file,
-                                                                        digest=digest)])
-
-            requests.append(remote_execution_pb2.BatchUpdateBlobsRequest.Request(
-                digest=digest, data=chunk))
-
-        request = remote_execution_pb2.BatchUpdateBlobsRequest(instance_name=instance_name,
-                                                               requests=requests)
+            command_arguments.append(argument.strip())
+
+        working_directory = None
+        if command.working_directory:
+            working_directory = os.path.join(temp_directory,
+                                             command.working_directory)
+            os.makedirs(working_directory, exist_ok=True)
+        else:
+            working_directory = temp_directory
+
+        # Ensure that output files structure exists:
+        for output_path in command.output_files:
+            directory_path = os.path.join(working_directory,
+                                          os.path.dirname(output_path))
+            os.makedirs(directory_path, exist_ok=True)
+
+        process = subprocess.Popen(command_arguments,
+                                   cwd=working_directory,
+                                   universal_newlines=True,
+                                   env=execution_envionment)
+        stdout_data, stderr_data = process.communicate()
+
+        update_requests = remote_execution_pb2.BatchUpdateBlobsRequest(instance_name=instance_name)
+        action_result = remote_execution_pb2.ActionResult()
+
+        for output_path in command.output_files:
+            file_path = os.path.join(working_directory, output_path)
+            # Missing outputs should simply be omitted in ActionResult:
+            if not os.path.isfile(file_path):
+                continue
+
+            output_file, update_request = output_file_maker(file_path, working_directory)
+
+            action_result.output_files.extend([output_file])
+            update_requests.requests.extend([update_request])
 
         stub_cas = remote_execution_pb2_grpc.ContentAddressableStorageStub(context.channel)
-        stub_cas.BatchUpdateBlobs(request)
+        stub_cas.BatchUpdateBlobs(update_requests)
 
         result_any = any_pb2.Any()
-        result_any.Pack(result)
+        result_any.Pack(action_result)
 
         lease.result.CopyFrom(result_any)