xref: /aosp_15_r20/external/pytorch/benchmarks/dynamo/test.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1import os
2import unittest
3
4from .common import parse_args, run
5from .torchbench import setup_torchbench_cwd, TorchBenchmarkRunner
6
7
8try:
9    # fbcode only
10    from aiplatform.utils.sanitizer_status import is_asan_or_tsan
11except ImportError:
12
13    def is_asan_or_tsan():
14        return False
15
16
17class TestDynamoBenchmark(unittest.TestCase):
18    @unittest.skipIf(is_asan_or_tsan(), "ASAN/TSAN not supported")
19    def test_benchmark_infra_runs(self) -> None:
20        """
21        Basic smoke test that TorchBench runs.
22
23        This test is mainly meant to check that our setup in fbcode
24        doesn't break.
25
26        If you see a failure here related to missing CPP headers, then
27        you likely need to update the resources list in:
28            //caffe2:inductor
29        """
30        original_dir = setup_torchbench_cwd()
31        try:
32            args = parse_args(
33                [
34                    "-dcpu",
35                    "--inductor",
36                    "--training",
37                    "--performance",
38                    "--only=BERT_pytorch",
39                    "-n1",
40                    "--batch-size=1",
41                ]
42            )
43            run(TorchBenchmarkRunner(), args, original_dir)
44        finally:
45            os.chdir(original_dir)
46