Add tools for generating images.

This formalizes the adhoc scripts previously in kokoro. The image targets can
be used by e.g. benchmarks in order to automated image prepation.

PiperOrigin-RevId: 290982744

1 files changed, 178 insertions, 0 deletions

diff --git a/tools/images/defs.bzl b/tools/images/defs.bzl
new file mode 100644
index 000000000..d8e422a5d
--- /dev/null
+++ b/tools/images/defs.bzl
@@ -0,0 +1,178 @@
+"""Image configuration.
+
+Images can be generated by using the vm_image rule. For example,
+
+  vm_image(
+      name = "ubuntu",
+      project = "...",
+      family = "...",
+      scripts = [
+          "script.sh",
+          "other.sh",
+      ],
+  )
+
+This will always create an vm_image in the current default gcloud project. The
+rule has a text file as its output containing the image name. This will enforce
+serialization for all dependent rules.
+
+Images are always named per the hash of all the hermetic input scripts. This
+allows images to be memoized quickly and easily.
+
+The vm_test rule can be used to execute a command remotely. For example,
+
+  vm_test(
+      name = "mycommand",
+      image = ":myimage",
+      targets = [":test"],
+  )
+"""
+
+def _vm_image_impl(ctx):
+    script_paths = []
+    for script in ctx.files.scripts:
+        script_paths.append(script.short_path)
+
+    resolved_inputs, argv, runfiles_manifests = ctx.resolve_command(
+        command = "USERNAME=%s ZONE=$(cat %s) IMAGE_PROJECT=%s IMAGE_FAMILY=%s %s %s > %s" %
+                  (
+                      ctx.attr.username,
+                      ctx.files.zone[0].path,
+                      ctx.attr.project,
+                      ctx.attr.family,
+                      ctx.executable.builder.path,
+                      " ".join(script_paths),
+                      ctx.outputs.out.path,
+                  ),
+        tools = [ctx.attr.builder] + ctx.attr.scripts,
+    )
+
+    ctx.actions.run_shell(
+        tools = resolved_inputs,
+        outputs = [ctx.outputs.out],
+        progress_message = "Building image...",
+        execution_requirements = {"local": "true"},
+        command = argv,
+        input_manifests = runfiles_manifests,
+    )
+    return [DefaultInfo(files = depset([ctx.outputs.out]))]
+
+_vm_image = rule(
+    attrs = {
+        "builder": attr.label(
+            executable = True,
+            default = "//tools/images:builder",
+            cfg = "host",
+        ),
+        "username": attr.string(default = "$(whoami)"),
+        "zone": attr.label(
+            default = "//tools/images:zone",
+            cfg = "host",
+        ),
+        "family": attr.string(mandatory = True),
+        "project": attr.string(mandatory = True),
+        "scripts": attr.label_list(allow_files = True),
+    },
+    outputs = {
+        "out": "%{name}.txt",
+    },
+    implementation = _vm_image_impl,
+)
+
+def vm_image(**kwargs):
+    _vm_image(
+        tags = [
+            "local",
+            "manual",
+        ],
+        **kwargs
+    )
+
+def _vm_test_impl(ctx):
+    runner = ctx.actions.declare_file("%s-executer" % ctx.label.name)
+
+    # Note that the remote execution case must actually generate an
+    # intermediate target in order to collect all the relevant runfiles so that
+    # they can be copied over for remote execution.
+    runner_content = "\n".join([
+        "#!/bin/bash",
+        "export ZONE=$(cat %s)" % ctx.files.zone[0].short_path,
+        "export USERNAME=%s" % ctx.attr.username,
+        "export IMAGE=$(cat %s)" % ctx.files.image[0].short_path,
+        "export SUDO=%s" % "true" if ctx.attr.sudo else "false",
+        "%s %s" % (
+            ctx.executable.executer.short_path,
+            " ".join([
+                target.files_to_run.executable.short_path
+                for target in ctx.attr.targets
+            ]),
+        ),
+        "",
+    ])
+    ctx.actions.write(runner, runner_content, is_executable = True)
+
+    # Return with all transitive files.
+    runfiles = ctx.runfiles(
+        transitive_files = depset(transitive = [
+            depset(target.data_runfiles.files)
+            for target in ctx.attr.targets
+            if hasattr(target, "data_runfiles")
+        ]),
+        files = ctx.files.executer + ctx.files.zone + ctx.files.image +
+                ctx.files.targets,
+        collect_default = True,
+        collect_data = True,
+    )
+    return [DefaultInfo(executable = runner, runfiles = runfiles)]
+
+_vm_test = rule(
+    attrs = {
+        "image": attr.label(
+            mandatory = True,
+            cfg = "host",
+        ),
+        "executer": attr.label(
+            executable = True,
+            default = "//tools/images:executer",
+            cfg = "host",
+        ),
+        "username": attr.string(default = "$(whoami)"),
+        "zone": attr.label(
+            default = "//tools/images:zone",
+            cfg = "host",
+        ),
+        "sudo": attr.bool(default = True),
+        "machine": attr.string(default = "n1-standard-1"),
+        "targets": attr.label_list(
+            mandatory = True,
+            allow_empty = False,
+            cfg = "target",
+        ),
+    },
+    test = True,
+    implementation = _vm_test_impl,
+)
+
+def vm_test(
+        installer = "//tools/installers:head",
+        **kwargs):
+    """Runs the given targets as a remote test.
+
+    Args:
+      installer: Script to run before all targets.
+      **kwargs: All test arguments. Should include targets and image.
+    """
+    targets = kwargs.pop("targets", [])
+    if installer:
+        targets = [installer] + targets
+    targets = [
+    ] + targets
+    _vm_test(
+        tags = [
+            "local",
+            "manual",
+        ],
+        targets = targets,
+        local = 1,
+        **kwargs
+    )