Merge #153167

153167: roachprod-microbench: support long retries when staging r=golgeek,DarrylWong a=herkolategan

Previously, sometimes while trying to stage the benchmark binaries the VMs might get preempted. But since this is a managed group and the VMs will come online again, we should retry for longer than the standard retry options.

This change adds support for longer retries when staging benchmark binaries. It is enabled by default since this binary is mostly used in the weekly job.

Epic: None
Release note: None

Co-authored-by: Herko Lategan <[email protected]>

Author: craig[bot]

pkg/cmd/roachprod-microbench/main.go

@@ -121,11 +121,12 @@ func makeRunCommand() *cobra.Command {
 }
 
 func makeStageCommand() *cobra.Command {
+	longRetries := true
 	runCmdFunc := func(cmd *cobra.Command, args []string) error {
 		cluster := args[0]
 		src := args[1]
 		dest := args[2]
-		return stage(cluster, src, dest)
+		return stage(cluster, src, dest, longRetries)
 	}
 
 	cmd := &cobra.Command{
@@ -139,6 +140,7 @@ The source can be a local path or a GCS URI.`,
 		RunE: runCmdFunc,
 	}
 	cmd.Flags().BoolVar(&roachprodConfig.Quiet, "quiet", roachprodConfig.Quiet, "suppress roachprod progress output")
+	cmd.Flags().BoolVar(&longRetries, "long-retries", longRetries, "use longer retry intervals to handle transient errors (for use with recoverable managed instances)")
 	return cmd
 }
 

pkg/cmd/roachprod-microbench/roachprod_util.go

@@ -23,11 +23,13 @@ func InitRoachprod() {
 
 // RoachprodRun runs a command on a roachprod cluster with the given cluster name and logger.
 // It takes a list of command arguments and passes them to the roachprod command execution.
-func RoachprodRun(clusterName string, l *logger.Logger, cmdArray []string) error {
+func RoachprodRun(
+	clusterName string, l *logger.Logger, cmdArray []string, runOptions install.RunOptions,
+) error {
 	// Execute the roachprod command with the provided context, logger, cluster name, and options.
 	return roachprod.Run(
 		context.Background(), l, clusterName, "", "", install.SimpleSecureOption(false),
-		os.Stdout, os.Stderr, cmdArray, install.DefaultRunOptions(),
+		os.Stdout, os.Stderr, cmdArray, runOptions,
 	)
 }
 

pkg/cmd/roachprod-microbench/stage.go

@@ -10,15 +10,18 @@ import (
 	"os"
 	"path"
 	"strings"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/roachprod"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
+	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/errors"
 )
 
 // stage copies the specified archive to the remote machine and extracts it to
 // the specified directory (creating it if it does not exist).
-func stage(cluster, archivePath, remoteDest string) (err error) {
+func stage(cluster, archivePath, remoteDest string, longRetries bool) (err error) {
 	ctx := context.Background()
 
 	InitRoachprod()
@@ -28,27 +31,41 @@ func stage(cluster, archivePath, remoteDest string) (err error) {
 		return
 	}
 
+	runOptions := install.DefaultRunOptions()
+	if longRetries {
+		// For VMs that have started failing with transient errors (usually due to
+		// preemption), we want to introduce a longer retry period. These VMs may
+		// still be recoverable since they're part of a managed instance group
+		// that attempts to recover failed VMs.
+		runOptions = runOptions.WithRetryOpts(retry.Options{
+			InitialBackoff: 1 * time.Minute,
+			MaxBackoff:     5 * time.Minute,
+			Multiplier:     2,
+			MaxRetries:     10,
+		})
+	}
+
 	archiveName := path.Base(archivePath)
 	archiveRemotePath := path.Join("/tmp", archiveName)
 
 	defer func() {
 		// Remove the remote archive after we're done.
-		cleanUpErr := RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath})
+		cleanUpErr := RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath}, runOptions)
 		err = errors.CombineErrors(err, errors.Wrapf(cleanUpErr, "removing remote archive: %s", archiveRemotePath))
 	}()
 
 	// Remove the remote archive and destination directory if they exist.
-	if err = RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath}, runOptions); err != nil {
 		return errors.Wrapf(err, "removing remote archive: %s", archiveRemotePath)
 	}
-	if err = RoachprodRun(cluster, l, []string{"rm", "-rf", remoteDest}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"rm", "-rf", remoteDest}, runOptions); err != nil {
 		return errors.Wrapf(err, "removing remote destination: %s", remoteDest)
 	}
 
 	// Copy the archive to the remote machine.
 	copyFromGCS := strings.HasPrefix(archivePath, "gs://")
 	if copyFromGCS {
-		if err = RoachprodRun(cluster, l, []string{"gsutil", "-q", "-m", "cp", archivePath, archiveRemotePath}); err != nil {
+		if err = RoachprodRun(cluster, l, []string{"gsutil", "-q", "-m", "cp", archivePath, archiveRemotePath}, runOptions); err != nil {
 			return errors.Wrapf(err, "copying archive from GCS: %s", archivePath)
 		}
 	} else {
@@ -58,10 +75,10 @@ func stage(cluster, archivePath, remoteDest string) (err error) {
 	}
 
 	// Extract the archive on the remote machine.
-	if err = RoachprodRun(cluster, l, []string{"mkdir", "-p", remoteDest}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"mkdir", "-p", remoteDest}, runOptions); err != nil {
 		return errors.Wrapf(err, "creating remote destination: %s", remoteDest)
 	}
-	if err = RoachprodRun(cluster, l, []string{"tar", "-C", remoteDest, "-xzf", archiveRemotePath}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"tar", "-C", remoteDest, "-xzf", archiveRemotePath}, runOptions); err != nil {
 		return errors.Wrapf(err, "extracting archive: %s", archiveRemotePath)
 	}