From 2db971f5266d473cba5bbfa9bc17c9f1aacc6869 Mon Sep 17 00:00:00 2001
From: Zach Koopmans <zkoopmans@google.com>
Date: Tue, 10 Jan 2023 11:10:13 -0800
Subject: [PATCH] Reduce number of Tensorflow test cases for dashboard.

A full run of all tensorflow test cases can take upwards of an hour
with very little benefit gained for most workloads. Reduce the number
of cases for dashboard runs.

PiperOrigin-RevId: 501043752
---
 .buildkite/pipeline.yaml              |  2 +-
 test/benchmarks/ml/tensorflow_test.go | 38 +++++++++++++++++----------
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/.buildkite/pipeline.yaml b/.buildkite/pipeline.yaml
index 29880f217..05a75514b 100644
--- a/.buildkite/pipeline.yaml
+++ b/.buildkite/pipeline.yaml
@@ -509,7 +509,7 @@ steps:
     command: make -i benchmark-platforms BENCHMARKS_SUITE=sysbench BENCHMARKS_TARGETS=test/benchmarks/base:sysbench_test
   - <<: *benchmarks
     label: ":tensorflow: TensorFlow benchmarks"
-    command: make -i benchmark-platforms BENCHMARKS_SUITE=tensorflow BENCHMARKS_TARGETS=test/benchmarks/ml:tensorflow_test BENCHMARKS_OPTIONS=-test.benchtime=1ns
+    command: make -i benchmark-platforms BENCHMARKS_SUITE=tensorflow BENCHMARKS_TARGETS=test/benchmarks/ml:tensorflow_test BENCHMARKS_FILTER=BenchmarkTensorflowDashboard
   - <<: *benchmarks
     label: ":gear: Syscall benchmarks"
     command: make -i benchmark-platforms BENCHMARKS_SUITE=syscall BENCHMARKS_TARGETS=test/benchmarks/base:syscallbench_test
diff --git a/test/benchmarks/ml/tensorflow_test.go b/test/benchmarks/ml/tensorflow_test.go
index ee0c946b8..58d4f7fd2 100644
--- a/test/benchmarks/ml/tensorflow_test.go
+++ b/test/benchmarks/ml/tensorflow_test.go
@@ -23,32 +23,42 @@ import (
 	"gvisor.dev/gvisor/test/benchmarks/tools"
 )
 
+func BenchmarkTensorflowDashboard(b *testing.B) {
+	workloads := map[string]string{
+		"ConvolutionalNetwork": "3_NeuralNetworks/convolutional_network.py",
+		"LogisticRegression":   "2_BasicModels/logistic_regression.py",
+		"NeuralNetwork":        "3_NeuralNetworks/neural_network.py",
+	}
+	doTensorflowTest(b, workloads)
+}
+
 // BenchmarkTensorflow runs workloads from a TensorFlow tutorial.
 // See: https://github.com/aymericdamien/TensorFlow-Examples
 func BenchmarkTensorflow(b *testing.B) {
-	workloads := []struct {
-		name, file string
-	}{
-		{"GradientDecisionTree", "2_BasicModels/gradient_boosted_decision_tree.py"},
-		{"Kmeans", "2_BasicModels/kmeans.py"},
-		{"LogisticRegression", "2_BasicModels/logistic_regression.py"},
-		{"NearestNeighbor", "2_BasicModels/nearest_neighbor.py"},
-		{"RandomForest", "2_BasicModels/random_forest.py"},
-		{"ConvolutionalNetwork", "3_NeuralNetworks/convolutional_network.py"},
-		{"MultilayerPerceptron", "3_NeuralNetworks/multilayer_perceptron.py"},
-		{"NeuralNetwork", "3_NeuralNetworks/neural_network.py"},
+	workloads := map[string]string{
+		"GradientDecisionTree": "2_BasicModels/gradient_boosted_decision_tree.py",
+		"Kmeans":               "2_BasicModels/kmeans.py",
+		"LogisticRegression":   "2_BasicModels/logistic_regression.py",
+		"NearestNeighbor":      "2_BasicModels/nearest_neighbor.py",
+		"RandomForest":         "2_BasicModels/random_forest.py",
+		"ConvolutionalNetwork": "3_NeuralNetworks/convolutional_network.py",
+		"MultilayerPerceptron": "3_NeuralNetworks/multilayer_perceptron.py",
+		"NeuralNetwork":        "3_NeuralNetworks/neural_network.py",
 	}
+	doTensorflowTest(b, workloads)
+}
 
+func doTensorflowTest(b *testing.B, workloads map[string]string) {
 	machine, err := harness.GetMachine()
 	if err != nil {
 		b.Fatalf("failed to get machine: %v", err)
 	}
 	defer machine.CleanUp()
 
-	for _, workload := range workloads {
+	for name, file := range workloads {
 		runName, err := tools.ParametersToName(tools.Parameter{
 			Name:  "operation",
-			Value: workload.name,
+			Value: name,
 		})
 		if err != nil {
 			b.Fatalf("Failed to parse param: %v", err)
@@ -73,7 +83,7 @@ func BenchmarkTensorflow(b *testing.B) {
 					Image:   "benchmarks/tensorflow",
 					Env:     []string{"PYTHONPATH=$PYTHONPATH:/TensorFlow-Examples/examples"},
 					WorkDir: "/TensorFlow-Examples/examples",
-				}, "python", workload.file); err != nil {
+				}, "python", file); err != nil {
 					b.Errorf("failed to run container: %v logs: %s", err, out)
 				}
 				b.StopTimer()