xref: /aosp_15_r20/external/skia/infra/bots/task_drivers/bazel_test_benchmark/bazel_test_benchmark.go (revision c8dee2aa9b3f27cf6c858bd81872bdeb2c07ed17)
1// Copyright 2023 Google LLC
2//
3// Use of this source code is governed by a BSD-style license that can be
4// found in the LICENSE file.
5
6// This task driver runs a single Bazel test target representing one or more benchmarks, or a Bazel
7// test suite consisting of multiple such targets, using the provided config (which is assumed to
8// be in //bazel/buildrc). Benchmark results are uploaded to a GCS bucket for ingestion by Perf.
9// This task driver handles any setup steps needed to run Bazel on our CI machines before running
10// the task, such as setting up logs and the Bazel cache.
11
12package main
13
14import (
15	"context"
16	"flag"
17	"fmt"
18	"path/filepath"
19
20	"cloud.google.com/go/storage"
21	"go.skia.org/infra/go/auth"
22	sk_exec "go.skia.org/infra/go/exec"
23	"go.skia.org/infra/go/gcs"
24	"go.skia.org/infra/go/gcs/gcsclient"
25	"go.skia.org/infra/go/skerr"
26	"go.skia.org/infra/task_driver/go/lib/auth_steps"
27	"go.skia.org/infra/task_driver/go/lib/bazel"
28	"go.skia.org/infra/task_driver/go/lib/os_steps"
29	"go.skia.org/infra/task_driver/go/td"
30	"go.skia.org/skia/infra/bots/task_drivers/common"
31	"google.golang.org/api/option"
32)
33
34var (
35	// Required properties for this task.
36	//
37	// We want the cache to be on a bigger disk than default. The root disk, where the home directory
38	// (and default Bazel cache) lives, is only 15 GB on our GCE VMs.
39	projectId = flag.String("project_id", "", "ID of the Google Cloud project.")
40	taskID    = flag.String("task_id", "", "ID of this task.")
41	taskName  = flag.String("task_name", "", "Name of the task.")
42	workdir   = flag.String("workdir", ".", "Working directory.")
43
44	gitCommit        = flag.String("git_commit", "", "The git hash to which the data should be associated.")
45	changelistID     = flag.String("changelist_id", "", "Should be non-empty only when run on the CQ.")
46	patchsetOrderStr = flag.String("patchset_order", "", "Should be non-zero only when run on the CQ.")
47
48	// Optional flags.
49	local  = flag.Bool("local", false, "True if running locally (as opposed to on the CI/CQ)")
50	output = flag.String("o", "", "If provided, dump a JSON blob of step data to the given file. Prints to stdout if '-' is given.")
51)
52
53func main() {
54	bazelFlags := common.MakeBazelFlags(common.MakeBazelFlagsOpts{
55		Label:                true,
56		Config:               true,
57		DeviceSpecificConfig: true,
58	})
59
60	// StartRun calls flag.Parse().
61	ctx := td.StartRun(projectId, taskID, taskName, output, local)
62	defer td.EndRun(ctx)
63
64	bazelFlags.Validate(ctx)
65
66	wd, err := os_steps.Abs(ctx, *workdir)
67	if err != nil {
68		td.Fatal(ctx, err)
69	}
70
71	opts := bazel.BazelOptions{
72		CachePath: *bazelFlags.CacheDir,
73	}
74	if err := bazel.EnsureBazelRCFile(ctx, opts); err != nil {
75		td.Fatal(ctx, err)
76	}
77
78	// Make an HTTP client with the required permissions to upload to the perf.skia.org GCS bucket.
79	httpClient, _, err := auth_steps.InitHttpClient(ctx, *local, auth.ScopeReadWrite, auth.ScopeUserinfoEmail)
80	if err != nil {
81		td.Fatal(ctx, skerr.Wrap(err))
82	}
83
84	// Make a GCS client to to upload to the perf.skia.org GCS bucket.
85	store, err := storage.NewClient(ctx, option.WithHTTPClient(httpClient))
86	if err != nil {
87		td.Fatal(ctx, skerr.Wrap(err))
88	}
89	gcsClient := gcsclient.New(store, common.PerfGCSBucketName)
90
91	if err := run(ctx, *bazelFlags.CacheDir, taskDriverArgs{
92		BenchmarkInfo: common.BenchmarkInfo{
93			GitCommit:     *gitCommit,
94			TaskName:      *taskName,
95			TaskID:        *taskID,
96			ChangelistID:  *changelistID,
97			PatchsetOrder: *patchsetOrderStr,
98		},
99		checkoutDir:               filepath.Join(wd, "skia"),
100		bazelLabel:                *bazelFlags.Label,
101		bazelConfig:               *bazelFlags.Config,
102		deviceSpecificBazelConfig: *bazelFlags.DeviceSpecificConfig,
103	}, gcsClient); err != nil {
104		td.Fatal(ctx, err)
105	}
106}
107
108// taskDriverArgs gathers the inputs to this task driver, and decouples the task driver's
109// entry-point function from the command line flags, which facilitates writing unit tests.
110type taskDriverArgs struct {
111	common.BenchmarkInfo
112	checkoutDir               string
113	bazelLabel                string
114	bazelConfig               string
115	deviceSpecificBazelConfig string
116}
117
118// run is the entrypoint of this task driver.
119func run(ctx context.Context, bazelCacheDir string, tdArgs taskDriverArgs, gcsClient gcs.GCSClient) error {
120	outputsZipPath, err := common.ValidateLabelAndReturnOutputsZipPath(tdArgs.checkoutDir, tdArgs.bazelLabel)
121	if err != nil {
122		return skerr.Wrap(err)
123	}
124
125	testArgs := common.ComputeBenchmarkTestRunnerCLIFlags(tdArgs.BenchmarkInfo)
126	if err := bazelTest(ctx, tdArgs.checkoutDir, tdArgs.bazelLabel, tdArgs.bazelConfig, tdArgs.deviceSpecificBazelConfig, testArgs); err != nil {
127		return skerr.Wrap(err)
128	}
129
130	if err := common.UploadToPerf(ctx, gcsClient, tdArgs.BenchmarkInfo, outputsZipPath); err != nil {
131		return skerr.Wrap(err)
132	}
133
134	if !*local {
135		if err := common.BazelCleanIfLowDiskSpace(ctx, bazelCacheDir, tdArgs.checkoutDir, "bazelisk"); err != nil {
136			return skerr.Wrap(err)
137		}
138	}
139
140	return nil
141}
142
143// bazelTest runs the test referenced by the given fully qualified Bazel label under the given
144// config.
145func bazelTest(ctx context.Context, checkoutDir, label, config, deviceSpecificConfig string, testArgs []string) error {
146	args := []string{"test",
147		label,
148		"--config=" + config,               // Should be defined in //bazel/buildrc.
149		"--config=" + deviceSpecificConfig, // Should be defined in //bazel/devicesrc.
150		"--test_output=errors",
151		"--jobs=100",
152	}
153	for _, testArg := range testArgs {
154		args = append(args, "--test_arg="+testArg)
155	}
156
157	return td.Do(ctx, td.Props(fmt.Sprintf("Test %s with config %s", label, config)), func(ctx context.Context) error {
158		runCmd := &sk_exec.Command{
159			Name:       "bazelisk",
160			Args:       args,
161			InheritEnv: true, // Makes sure bazelisk is on PATH.
162			Dir:        checkoutDir,
163			LogStdout:  true,
164			LogStderr:  true,
165		}
166		_, err := sk_exec.RunCommand(ctx, runCmd)
167		if err != nil {
168			return err
169		}
170		return nil
171	})
172}
173