roachprod: pre-bake custom roachprod cloud images

Prior to this patch, roachprod clusters were created from bare Ubuntu
images.

This was inadequate for multiple reasons, some of which being:
- dependency on third-parties (GCS, APT repositories) availability
- spinning up two clusters at a different moment in time could lead to
  different resulting systems (package versions, ...) and create
reproducibility issues
- growing number of dependencies installed increases the boot time

To address this, this patch creates a new roachprod bake-images command
that relies on Hashicorp Packer to pre-bake ready to use cloud images
for AWS and GCP. This creates a system dependency on Packer and requires
the machine that runs the command to have Packer installed and to be
authenticated on AWS and GCP with authorization to create instances and
publish new images. If an image already exist, it won't get built again,
making re-running roachprod bake-images safe.

The pre-baking process creates images for amd64, arm64 and fips, and
pushes them to the roachprod compatible regions (only for AWS, since
images are globally available in GCP). The images are tagged with a
hashed checksum of the startup script, which defines their unique
version.

At runtime, the providers checksums the startup script to figure out
which pre-baked image should be used, and checks for its availability in
the cloud provider for that specific region/zone:
- if the image exists, it is used to create the instance, and only a
  subset (runtime) of the startup scripts is executed on the instances,
  decreasing the startup time to a minimum (5s or so for disk setup)
- if the image does not exists, the system fallbacks to using the base
  image and the whole startup scripts (pre-baking + runtime) is executed
  on the instances

This patch also drops the JSON hardcoded AMI IDs (or names in GCP) and
introduces auto-discovery of the base image's most recent version based
on the image name/family and owner or project ID. This allows us to
automatically keep up to date with the latest patch releases, which
usually are security updates.

Notes:
- this patch only contains implementation for AWS and GCP, and Azure and
  IBM should also be implemented
- a CI mechanism should be built to automatically build all images when
  there is a change in the startup scripts (either Github upon merge to
master or TeamCity nightly runs)
- there is currently no built-in way to deprecate/cleanup previous
  images since they might still be used on older branches; a cleanup
  routine should be considered if/when the number of images get out of
  hand

Beyond this first iteration, a concept of "pre-bake only snippets"
should come next: snippets that are only executed at pre-baking time and
not at runtime even if there is no pre-baked image.
These snippets would contain adhoc roachtest setups
(building/pre-installing third party tools like Prometheus/Grafana,
Jepsen, Kafka CLI, ...), which would remove the need for these tests to
build/install at third party dependencies at runtime if the test is
running on an instance supported by a pre-baked image (see #62066 as an
example).

Epic: none
Informs: #150144
Release note: None

Author: golgeek

pkg/BUILD.bazel

@@ -1743,6 +1743,7 @@ GO_TARGETS = [
     "//pkg/roachprod/vm/ibm:ibm_test",
     "//pkg/roachprod/vm/local:local",
     "//pkg/roachprod/vm/local:local_test",
+    "//pkg/roachprod/vm/utils/packer:packer",
     "//pkg/roachprod/vm:vm",
     "//pkg/roachprod/vm:vm_test",
     "//pkg/roachprod:roachprod",

pkg/cmd/roachprod/cli/BUILD.bazel

@@ -27,14 +27,17 @@ go_library(
         "//pkg/roachprod/ssh",
         "//pkg/roachprod/ui",
         "//pkg/roachprod/vm",
+        "//pkg/roachprod/vm/aws",
         "//pkg/roachprod/vm/gce",
+        "//pkg/roachprod/vm/utils/packer",
         "//pkg/util/envutil",
         "//pkg/util/flagutil",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_errors//oserror",
         "@com_github_fatih_color//:color",
         "@com_github_spf13_cobra//:cobra",
+        "@com_github_spf13_pflag//:pflag",
         "@com_google_cloud_go_storage//:storage",
         "@org_golang_google_api//option",
         "@org_golang_x_crypto//ssh",

pkg/cmd/roachprod/cli/commands.go

@@ -28,7 +28,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachprod/roachprodutil"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/ui"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/vm"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/vm/aws"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/vm/gce"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/vm/utils/packer"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"github.com/fatih/color"
@@ -2256,3 +2258,132 @@ If a destination is not provided, the certs will be downloaded to a default %s d
 		}),
 	}
 }
+
+func (cr *commandRegistry) buildBakeImageCmd() *cobra.Command {
+	var providers []string
+	var zones []string
+	var project string
+	var architectures []string
+	var dryRun bool
+
+	bakeImageCmd := &cobra.Command{
+		Use:   "bake-images",
+		Short: "pre-bake cloud VM images with roachprod configuration",
+		Long: `Pre-bake cloud VM images with packages and configuration to speed up cluster creation.
+
+This command creates cloud images with all roachprod packages (node_exporter,
+ebpf_exporter, chrony, etc.) pre-installed. When creating clusters with these
+pre-baked images, instance startup time is significantly reduced since packages
+don't need to be downloaded and installed.
+
+The image name is generated based on the base Ubuntu image and a checksum of
+the startup template content. If an image with the computed name already exists,
+the command exits successfully without creating a new image.
+
+If no provider is specified, images will be built for all supported providers.
+If no architecture is specified, all supported architectures will be built.
+
+Examples:
+  roachprod bake-image                                           # Build for all providers and architectures
+  roachprod bake-image --provider=gce --zones=us-central1-a      # GCE only
+  roachprod bake-image --provider=aws                            # AWS only
+  roachprod bake-image --provider=gce,aws --arch=amd64,arm64     # Both providers, specific architectures`,
+		Args: cobra.NoArgs,
+		Run: Wrap(func(cmd *cobra.Command, args []string) error {
+			// Default to all providers if none specified
+			if len(providers) == 0 {
+				providers = []string{"gce", "aws"}
+			}
+
+			// Default to all architectures if none specified
+			if len(architectures) == 0 {
+				architectures = []string{"amd64", "arm64", "fips"}
+			}
+
+			// Collect sources and provisioners from all providers
+			var allSources []packer.SourceConfig
+			var allProvisioners []packer.ProvisionerConfig
+			var allPlugins []packer.PluginConfig
+
+			for _, provider := range providers {
+				config.Logger.Printf("Preparing images for provider: %s", provider)
+
+				var sources []packer.SourceConfig
+				var provisioners []packer.ProvisionerConfig
+				var plugins []packer.PluginConfig
+				var err error
+
+				switch provider {
+				case "gce":
+					// Get GCE provider instance
+					gceProvider, ok := vm.Providers["gce"].(*gce.Provider)
+					if !ok {
+						return errors.New("GCE provider not initialized")
+					}
+
+					providerOpts := map[string]interface{}{
+						"zones":   zones,
+						"project": project,
+					}
+					sources, provisioners, plugins, err = gceProvider.GetPackerSources(
+						config.Logger, architectures, providerOpts,
+					)
+
+				case "aws":
+					// Get AWS provider instance
+					awsProvider, ok := vm.Providers["aws"].(*aws.Provider)
+					if !ok {
+						return errors.New("AWS provider not initialized")
+					}
+
+					providerOpts := map[string]interface{}{}
+					sources, provisioners, plugins, err = awsProvider.GetPackerSources(
+						config.Logger, architectures, providerOpts,
+					)
+
+				default:
+					return errors.Newf("unsupported provider: %s (supported: gce, aws)", provider)
+				}
+
+				if err != nil {
+					return errors.Wrapf(err, "failed to get Packer sources for %s", provider)
+				}
+
+				// Accumulate sources, provisioners, and plugins
+				allSources = append(allSources, sources...)
+				allProvisioners = append(allProvisioners, provisioners...)
+				allPlugins = append(allPlugins, plugins...)
+			}
+
+			// If no sources to build, we're done
+			if len(allSources) == 0 {
+				config.Logger.Printf("All requested images already exist, nothing to build")
+				return nil
+			}
+
+			// Build all images in a single Packer run (parallel across providers!)
+			config.Logger.Printf("Building images across %d provider(s) in parallel...", len(providers))
+			if err := packer.Build(config.Logger, allSources, allProvisioners, allPlugins, dryRun); err != nil {
+				return errors.Wrap(err, "packer build failed")
+			}
+
+			config.Logger.Printf("Successfully baked images for all requested providers")
+			return nil
+		}),
+	}
+
+	bakeImageCmd.Flags().StringSliceVar(&providers, "provider", nil,
+		"cloud provider(s) to build images for (gce, aws); if not specified, builds for all")
+	bakeImageCmd.Flags().StringSliceVar(&zones, "zones", []string{"us-central1-a"},
+		"zones to build the image in (GCE only, uses first zone)")
+	bakeImageCmd.Flags().StringVar(&project, "project", gce.DefaultProject(),
+		"GCE project to create the image in (GCE only)")
+	bakeImageCmd.Flags().StringSliceVar(&architectures, "arch", nil,
+		"architectures to build for (amd64, arm64, fips); if not specified, builds all")
+	bakeImageCmd.Flags().BoolVar(&dryRun, "dry-run", false,
+		"if set, only prints the actions without executing them")
+
+	cr.addToExcludeFromBashCompletion(bakeImageCmd)
+	cr.addToExcludeFromClusterFlagsMulti(bakeImageCmd)
+	return bakeImageCmd
+}

pkg/cmd/roachprod/cli/resgistry.go

@@ -73,5 +73,6 @@ func (cr *commandRegistry) register() {
 		cr.buildFetchLogsCmd(),
 		cr.buildGetLatestPProfCmd(),
 		cr.buildFetchCertsDir(),
+		cr.buildBakeImageCmd(),
 	})
 }

pkg/cmd/roachprod/cli/util.go

@@ -8,6 +8,7 @@ package cli
 import (
 	"fmt"
 	"os"
+	"slices"
 	"strings"
 	"text/tabwriter"
 	"time"
@@ -19,6 +20,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/errors/oserror"
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 )
 
 func PromptYesNo(msg string, defaultYes bool) bool {
@@ -211,16 +213,40 @@ func ValidateAndConfigure(cmd *cobra.Command, args []string) {
 		}
 	}
 
+	validArchitectures := []vm.CPUArch{vm.ArchAMD64, vm.ArchARM64, vm.ArchFIPS, vm.ArchS390x}
+
 	// Validate architecture flag, if set.
 	if archOpt := cmd.Flags().Lookup("arch"); archOpt != nil && archOpt.Changed {
-		arch := vm.CPUArch(strings.ToLower(archOpt.Value.String()))
 
-		if arch != vm.ArchAMD64 && arch != vm.ArchARM64 && arch != vm.ArchFIPS && arch != vm.ArchS390x {
-			printErrAndExit(fmt.Errorf("unsupported architecture %q", arch))
+		var architecturesToValidate = []string{}
+		switch v := archOpt.Value.(type) {
+		case pflag.SliceValue:
+			architecturesToValidate = v.GetSlice()
+		default:
+			architecturesToValidate = []string{archOpt.Value.String()}
+		}
+
+		normalizedArchitectures := make([]string, 0)
+		for _, archStr := range architecturesToValidate {
+
+			arch := vm.CPUArch(strings.ToLower(archStr))
+
+			if !slices.Contains(validArchitectures, arch) {
+				printErrAndExit(fmt.Errorf("unsupported architecture %q", archStr))
+			}
+
+			// Canonicalize architecture flag value.
+			normalizedArchitectures = append(normalizedArchitectures, string(arch))
 		}
-		if string(arch) != archOpt.Value.String() {
-			// Set the canonical value.
-			_ = cmd.Flags().Set("arch", string(arch))
+
+		// Replace the value by accessing the underlying slice directly
+		// This is a bit hacky but works with pflag
+		if sliceValue, ok := archOpt.Value.(pflag.SliceValue); ok {
+			// For StringSliceVar, replace the entire slice
+			_ = sliceValue.Replace(normalizedArchitectures)
+		} else {
+			// For StringVar, just set the first value
+			_ = cmd.Flags().Set("arch", normalizedArchitectures[0])
 		}
 	}
 

pkg/roachprod/vm/aws/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "aws.go",
         "config.go",
         "keys.go",
+        "pre_bake_amis.go",
         "support.go",
     ],
     embedsrcs = [
@@ -19,6 +20,7 @@ go_library(
         "//pkg/roachprod/logger",
         "//pkg/roachprod/vm",
         "//pkg/roachprod/vm/flagstub",
+        "//pkg/roachprod/vm/utils/packer",
         "//pkg/util/retry",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",