Merge pull request #2036 from apache/abderrahim/digest-environment

Add the digest-environment dependency configuration

Author: juergbi

src/buildstream/_artifact.py

@@ -200,6 +200,7 @@ def get_extract_key(self):
     #    variables (Variables): The element's Variables
     #    environment (dict): dict of the element's environment variables
     #    sandboxconfig (SandboxConfig): The element's SandboxConfig
+    #    buildsandbox (Sandbox): The element's configured build sandbox
     #
     def cache(
         self,
@@ -213,6 +214,7 @@ def cache(
         variables,
         environment,
         sandboxconfig,
+        buildsandbox,
     ):
 
         context = self._context
@@ -317,6 +319,19 @@ def cache(
             rootvdir._import_files_internal(buildrootvdir, properties=properties, collect_result=False)
             artifact.buildroot.CopyFrom(rootvdir._get_digest())
 
+        if buildsandbox is not None:
+            sandbox_env = buildsandbox._get_configured_environment()
+            if sandbox_env:
+                for key, value in sorted(sandbox_env.items()):
+                    artifact.buildsandbox.environment.add(name=key, value=value)
+
+            artifact.buildsandbox.working_directory = buildsandbox._get_work_directory()
+
+            for subsandbox in buildsandbox._get_subsandboxes():
+                vdir = subsandbox.get_virtual_directory()
+                digest = artifact.buildsandbox.subsandbox_digests.add()
+                digest.CopyFrom(vdir._get_digest())
+
         os.makedirs(os.path.dirname(os.path.join(self._artifactdir, element.get_artifact_name())), exist_ok=True)
         keys = utils._deduplicate([self._cache_key, self._weak_cache_key])
         for key in keys:
@@ -681,6 +696,21 @@ def pull(self, *, pull_buildtrees):
 
         return True
 
+    def configure_sandbox(self, sandbox):
+        artifact = self._get_proto()
+
+        if artifact.buildsandbox and artifact.buildsandbox.environment:
+            env = {}
+            for env_var in artifact.buildsandbox.environment:
+                env[env_var.name] = env_var.value
+        else:
+            env = self.load_environment()
+
+        sandbox.set_environment(env)
+
+        if artifact.buildsandbox and artifact.buildsandbox.working_directory:
+            sandbox.set_work_directory(artifact.buildsandbox.working_directory)
+
     #  load_proto()
     #
     # Returns:

src/buildstream/_artifactcache.py

@@ -303,6 +303,10 @@ def _push_artifact_blobs(self, artifact, artifact_digest, remote):
                 except FileNotFoundError:
                     pass
 
+            if artifact_proto.buildsandbox:
+                for subsandbox_digest in artifact_proto.buildsandbox.subsandbox_digests:
+                    self.cas._send_directory(remote, subsandbox_digest)
+
             digests = [artifact_digest, artifact_proto.low_diversity_meta, artifact_proto.high_diversity_meta]
 
             if str(artifact_proto.public_data):
@@ -361,6 +365,9 @@ def _push_artifact_proto(self, element, artifact, artifact_digest, remote):
             referenced_directories.append(artifact_proto.sources)
         if artifact_proto.buildroot:
             referenced_directories.append(artifact_proto.buildroot)
+        if artifact_proto.buildsandbox:
+            for subsandbox_digest in artifact_proto.buildsandbox.subsandbox_digests:
+                referenced_directories.append(subsandbox_digest)
 
         referenced_blobs = [artifact_proto.low_diversity_meta, artifact_proto.high_diversity_meta] + [
             log_file.digest for log_file in artifact_proto.logs
@@ -419,6 +426,9 @@ def _pull_artifact_storage(self, element, key, artifact_digest, remote, pull_bui
                     self.cas.fetch_directory(remote, artifact.buildtree)
                 if str(artifact.buildroot):
                     self.cas.fetch_directory(remote, artifact.buildroot)
+                if artifact.buildsandbox:
+                    for subsandbox_digest in artifact.buildsandbox.subsandbox_digests:
+                        self.cas.fetch_directory(remote, subsandbox_digest)
 
             digests = [artifact.low_diversity_meta, artifact.high_diversity_meta]
             if str(artifact.public_data):

src/buildstream/_elementproxy.py

@@ -104,9 +104,10 @@ def stage_artifact(
         owner = cast("Element", self._owner)
         element = cast("Element", self._plugin)
 
-        assert owner._overlap_collector is not None, "Attempted to stage artifacts outside of Element.stage()"
+        overlap_collector = owner._overlap_collectors.get(sandbox)
+        assert overlap_collector is not None, "Attempted to stage artifacts outside of Element.stage()"
 
-        with owner._overlap_collector.session(action, path):
+        with overlap_collector.session(action, path):
             result = element._stage_artifact(
                 sandbox, path=path, action=action, include=include, exclude=exclude, orphans=orphans, owner=owner
             )

src/buildstream/_protos/buildstream/v2/artifact.proto

@@ -85,4 +85,11 @@ message Artifact {
 
   // digest of a directory
   build.bazel.remote.execution.v2.Digest buildroot = 17;  // optional
+
+  message SandboxState {
+    repeated build.bazel.remote.execution.v2.Command.EnvironmentVariable environment = 1;
+    string working_directory = 2;
+    repeated build.bazel.remote.execution.v2.Digest subsandbox_digests = 3;
+  };
+  SandboxState buildsandbox = 18;  // optional
 }

src/buildstream/_protos/buildstream/v2/artifact_pb2.py

@@ -8,7 +8,7 @@ from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Map
 DESCRIPTOR: _descriptor.FileDescriptor
 
 class Artifact(_message.Message):
-    __slots__ = ("version", "build_success", "build_error", "build_error_details", "strong_key", "weak_key", "was_workspaced", "files", "build_deps", "public_data", "logs", "buildtree", "sources", "low_diversity_meta", "high_diversity_meta", "strict_key", "buildroot")
+    __slots__ = ("version", "build_success", "build_error", "build_error_details", "strong_key", "weak_key", "was_workspaced", "files", "build_deps", "public_data", "logs", "buildtree", "sources", "low_diversity_meta", "high_diversity_meta", "strict_key", "buildroot", "buildsandbox")
     class Dependency(_message.Message):
         __slots__ = ("project_name", "element_name", "cache_key", "was_workspaced")
         PROJECT_NAME_FIELD_NUMBER: _ClassVar[int]
@@ -27,6 +27,15 @@ class Artifact(_message.Message):
         name: str
         digest: _remote_execution_pb2.Digest
         def __init__(self, name: _Optional[str] = ..., digest: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ...) -> None: ...
+    class SandboxState(_message.Message):
+        __slots__ = ("environment", "working_directory", "subsandbox_digests")
+        ENVIRONMENT_FIELD_NUMBER: _ClassVar[int]
+        WORKING_DIRECTORY_FIELD_NUMBER: _ClassVar[int]
+        SUBSANDBOX_DIGESTS_FIELD_NUMBER: _ClassVar[int]
+        environment: _containers.RepeatedCompositeFieldContainer[_remote_execution_pb2.Command.EnvironmentVariable]
+        working_directory: str
+        subsandbox_digests: _containers.RepeatedCompositeFieldContainer[_remote_execution_pb2.Digest]
+        def __init__(self, environment: _Optional[_Iterable[_Union[_remote_execution_pb2.Command.EnvironmentVariable, _Mapping]]] = ..., working_directory: _Optional[str] = ..., subsandbox_digests: _Optional[_Iterable[_Union[_remote_execution_pb2.Digest, _Mapping]]] = ...) -> None: ...
     VERSION_FIELD_NUMBER: _ClassVar[int]
     BUILD_SUCCESS_FIELD_NUMBER: _ClassVar[int]
     BUILD_ERROR_FIELD_NUMBER: _ClassVar[int]
@@ -44,6 +53,7 @@ class Artifact(_message.Message):
     HIGH_DIVERSITY_META_FIELD_NUMBER: _ClassVar[int]
     STRICT_KEY_FIELD_NUMBER: _ClassVar[int]
     BUILDROOT_FIELD_NUMBER: _ClassVar[int]
+    BUILDSANDBOX_FIELD_NUMBER: _ClassVar[int]
     version: int
     build_success: bool
     build_error: str
@@ -61,4 +71,5 @@ class Artifact(_message.Message):
     high_diversity_meta: _remote_execution_pb2.Digest
     strict_key: str
     buildroot: _remote_execution_pb2.Digest
-    def __init__(self, version: _Optional[int] = ..., build_success: bool = ..., build_error: _Optional[str] = ..., build_error_details: _Optional[str] = ..., strong_key: _Optional[str] = ..., weak_key: _Optional[str] = ..., was_workspaced: bool = ..., files: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., build_deps: _Optional[_Iterable[_Union[Artifact.Dependency, _Mapping]]] = ..., public_data: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., logs: _Optional[_Iterable[_Union[Artifact.LogFile, _Mapping]]] = ..., buildtree: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., sources: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., low_diversity_meta: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., high_diversity_meta: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., strict_key: _Optional[str] = ..., buildroot: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ...) -> None: ...
+    buildsandbox: Artifact.SandboxState
+    def __init__(self, version: _Optional[int] = ..., build_success: bool = ..., build_error: _Optional[str] = ..., build_error_details: _Optional[str] = ..., strong_key: _Optional[str] = ..., weak_key: _Optional[str] = ..., was_workspaced: bool = ..., files: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., build_deps: _Optional[_Iterable[_Union[Artifact.Dependency, _Mapping]]] = ..., public_data: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., logs: _Optional[_Iterable[_Union[Artifact.LogFile, _Mapping]]] = ..., buildtree: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., sources: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., low_diversity_meta: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., high_diversity_meta: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., strict_key: _Optional[str] = ..., buildroot: _Optional[_Union[_remote_execution_pb2.Digest, _Mapping]] = ..., buildsandbox: _Optional[_Union[Artifact.SandboxState, _Mapping]] = ...) -> None: ...

src/buildstream/_protos/buildstream/v2/artifact_pb2.pyi

@@ -76,6 +76,31 @@
     directories before subdirectories.
 
 
+`digest-environment` for dependencies
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The BuildElement supports the ``digest-environment`` :term:`dependency configuration <Dependency configuration>`,
+which sets the specified environment variable in the build sandbox to the CAS digest
+corresponding to a directory that contains all dependencies that are configured
+with the same ``digest-environment``.
+
+This is useful for REAPI clients in the sandbox such as `recc <https://buildgrid.gitlab.io/recc>`_,
+see ``remote-apis-socket`` in the :ref:`sandbox configuration <format_sandbox>`.
+
+**Example:**
+
+Here is an example of how to set the environment variable `GCC_DIGEST` to the
+CAS digest of a directory that contains ``gcc.bst`` and its runtime dependencies.
+The ``libpony.bst`` dependency will not be included in that CAS directory.
+
+.. code:: yaml
+
+   build-depends:
+   - baseproject.bst:gcc.bst
+     config:
+       digest-environment: GCC_DIGEST
+   - libpony.bst
+
+
 Location for running commands
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The ``command-subdir`` variable sets where commands will be executed,
@@ -219,6 +244,7 @@ def configure(self, node):
     def configure_dependencies(self, dependencies):
 
         self.__layout = {}  # pylint: disable=attribute-defined-outside-init
+        self.__digest_environment = {}  # pylint: disable=attribute-defined-outside-init
 
         # FIXME: Currently this forcefully validates configurations
         #        for all BuildElement subclasses so they are unable to
@@ -227,9 +253,18 @@ def configure_dependencies(self, dependencies):
         for dep in dependencies:
             # Determine the location to stage each element, default is "/"
             location = "/"
+
             if dep.config:
-                dep.config.validate_keys(["location"])
-                location = dep.config.get_str("location")
+                dep.config.validate_keys(["digest-environment", "location"])
+
+                location = dep.config.get_str("location", "/")
+
+                digest_var_name = dep.config.get_str("digest-environment", None)
+
+                if digest_var_name is not None:
+                    element_list = self.__digest_environment.setdefault(digest_var_name, [])
+                    element_list.append((dep.element, dep.path))
+
             try:
                 element_list = self.__layout[location]
             except KeyError:
@@ -268,6 +303,16 @@ def get_unique_key(self):
             }
             dictionary["layout"] = layout_key
 
+        # Specify the layout in the key, if buildstream is to generate an environment
+        # variable with the digest
+        #
+        if self.__digest_environment:
+            sorted_envs = sorted(self.__digest_environment)
+            digest_key = {
+                env: [dependency_path for _, dependency_path in self.__digest_environment[env]] for env in sorted_envs
+            }
+            dictionary["digest-enviornment"] = digest_key
+
         return dictionary
 
     def configure_sandbox(self, sandbox):
@@ -286,7 +331,20 @@ def configure_sandbox(self, sandbox):
         sandbox.set_work_directory(command_dir)
 
         # Setup environment
-        sandbox.set_environment(self.get_environment())
+        env = self.get_environment()
+
+        # Add "CAS digest" environment variables
+        sorted_envs = sorted(self.__digest_environment)
+        for digest_variable in sorted_envs:
+            element_list = [element for element, _ in self.__digest_environment[digest_variable]]
+            with self.timed_activity(
+                f"Staging dependencies for '{digest_variable}' in subsandbox", silent_nested=True
+            ), self.subsandbox(sandbox) as subsandbox:
+                self.stage_dependency_artifacts(subsandbox, element_list)
+                digest = subsandbox.get_virtual_directory()._get_digest()
+            env[digest_variable] = "{}/{}".format(digest.hash, digest.size_bytes)
+
+        sandbox.set_environment(env)
 
     def stage(self, sandbox):