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.

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

Epic: none
Informs: #150144
Release note: None

Author: golgeek

pkg/BUILD.bazel

@@ -1739,6 +1739,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,7 +27,9 @@ 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",

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,127 @@ 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
+
+	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); 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")
+
+	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/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",