Add a new ai benchmark suite in DCPerf (#139)

Summary:
Pull Request resolved: #139

This diff introduces a benchmark suite named `ai` into DCPerf. The `ai` benchmark suite currently includes two embedding jobs, which will install `fbgemm_gpu` and utilize the corresponding benchmark with representative parameters collected from a production model.

To use it,
```
./benchpress -b ai list
./benchpress -b ai install
./benchpress -b ai run
```

The changes include:
- Add a new parser for the embedding parser.
- Add the `benchmarks_ai.yml` for the ai benchmark suite.
- Update the `jobs_ai.yml` file to include jobs for the ai benchmark suite.
- Update the `__init__.py `file to register the new parser and new benchmark suite.
- `install_embedding.sh` builds the fbgemm_gpu from source.
- `cleanup_embedding.sh` clean up the exec from benchmarks directory.

Reviewed By: excelle08

Differential Revision: D76377213

fbshipit-source-id: a97a0b35703365f07aac69ea19ee53874c585e5e

Author: Kaiwei Tu

benchpress/config/__init__.py

@@ -70,6 +70,7 @@ def register_benchmark_suite(name):
     register_benchmark_suite("internal")
 register_benchmark_suite("wdl")
 register_benchmark_suite("system")
+register_benchmark_suite("ai")
 
 
 class BenchpressConfig:

benchpress/config/benchmarks_ai.yml

@@ -0,0 +1,7 @@
+fbgemm_embedding:
+  parser: embedding
+  install_script: ./packages/fbgemm_gpu/install_embedding.sh
+  cleanup_script: ./packages/fbgemm_gpu/cleanup_embedding.sh
+  path: ./benchmarks/fbgemm_embedding/tbe_inference_benchmark
+  metrics:
+    - bandwidth

benchpress/config/jobs_ai.yml

@@ -0,0 +1,70 @@
+- benchmark: fbgemm_embedding
+  name: fbgemm_embedding_a_single
+  description: Embedding workload for ai benchmark with one representative table
+  args:
+    - nbit-cpu
+    - '--num-embeddings={embeddings}'
+    - '--bag-size={bag_size}'
+    - '--embedding-dim={embedding_dim}'
+    - '--batch-size={batch_size}'
+    - '--num-tables={num_tables}'
+    - '--weights-precision={weights_precision}'
+    - '--copies={copies}'
+    - '--iters={iters}'
+
+  vars:
+    - 'embeddings=40000000'
+    - 'bag_size=2'
+    - 'embedding_dim=96'
+    - 'batch_size=166'
+    - 'num_tables=1'
+    - 'weights_precision=int4'
+    - 'copies=16'
+    - 'iters=1000'
+
+- benchmark: fbgemm_embedding
+  name: fbgemm_embedding_a_spec
+  description: Embedding workload for ai benchmark with 8 representative tables having different bag sizes.
+  args:
+    - nbit-device-with-spec
+    - '--num-embeddings-list={embeddings_list}'
+    - '--bag-size-list={bag_size_list}'
+    - '--embedding-dim-list={embedding_dim_list}'
+    - '--batch-size={batch_size}'
+    - '--weights-precision={weights_precision}'
+    - '--cpu-copies={copies}'
+    - '--iters={iters}'
+    - '--use-cpu'
+
+  vars:
+    - 'embeddings_list=40000000,40000000,40000000,40000000,40000000,40000000,40000000,40000000'
+    - 'bag_size_list=1,1,1,1,1,1,1,5'
+    - 'embedding_dim_list=96,96,96,96,96,96,96,96'
+    - 'batch_size=166'
+    - 'weights_precision=int4'
+    - 'copies=16'
+    - 'iters=1000'
+
+
+- benchmark: fbgemm_embedding
+  name: fbgemm_embedding_b_spec
+  description: Embedding workload for ai benchmark with 8 representative tables having different bag sizes.
+  args:
+    - nbit-device-with-spec
+    - '--num-embeddings-list={embeddings_list}'
+    - '--bag-size-list={bag_size_list}'
+    - '--embedding-dim-list={embedding_dim_list}'
+    - '--batch-size={batch_size}'
+    - '--weights-precision={weights_precision}'
+    - '--cpu-copies={copies}'
+    - '--iters={iters}'
+    - '--use-cpu'
+
+  vars:
+    - 'embeddings_list=100000,100000,100000,100000,100000,5000000,5000000,5000000,10000,10000,14794452'
+    - 'bag_size_list=1,1,1,1,1,1,1,1,1,1,1'
+    - 'embedding_dim_list=64,64,64,64,64,64,64,64,64,64,64'
+    - 'batch_size=166'
+    - 'weights_precision=int4'
+    - 'copies=16'
+    - 'iters=1000'

benchpress/plugins/parsers/__init__.py

@@ -15,6 +15,7 @@
 from .cloudsuite_graph import CloudSuiteGraphParser
 from .compression_parser import CompressionParser
 from .django_workload import DjangoWorkloadParser
+from .embedding import EmbeddingParser
 from .encryption import EncryptionParser
 from .fb_fiosynth import Fiosynth_Parser
 from .fbgemm import FbgemmParser
@@ -100,6 +101,7 @@ def register_parsers(factory):
     factory.register("wdl_bench", WDLParser)
     factory.register("health_check", HealthCheckParser)
     factory.register("syscall", SyscallParser)
+    factory.register("embedding", EmbeddingParser)
 
     if not open_source:
         factory.register("adsim", AdSimParser)

benchpress/plugins/parsers/embedding.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+import re
+
+from benchpress.lib.parser import Parser
+
+
+class EmbeddingParser(Parser):
+    def parse(self, stdout, stderr, returncode):
+        metrics = {}
+        for line in stderr:
+            match = re.search(r"BW:\s*(\d+\.\d+)\s*GB/s", line)
+            if match:
+                metrics["bandwidth"] = float(match.group(1))
+
+        return metrics

packages/fbgemm_gpu/cleanup_embedding.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+AI_BENCH_ROOT="$(dirname "$(readlink -f "$0")")" # Path to dir with this file.
+BENCHPRESS_ROOT="$(readlink -f "$AI_BENCH_ROOT/../..")"
+BENCHMARKS_DIR="${BENCHPRESS_ROOT}/benchmarks/fbgemm_embedding"
+MEMCACHE_BENCH_DIR="${BENCHMARKS_DIR}/tbe_inference_benchmark"
+
+rm -rf "$MEMCACHE_BENCH_DIR"