From 181f91d1a091f33bd52f8a70ba698d20e15bb014 Mon Sep 17 00:00:00 2001
From: Yukio Siraichi <yukio.siraichi@gmail.com>
Date: Wed, 13 Dec 2023 02:33:48 -0300
Subject: [PATCH 1/2] Move tests to xla/test/benchmark.

---
 .circleci/common.sh                                  |  5 -----
 benchmarks/test/run_tests.sh                         | 12 ------------
 .../benchmarks/output-example.json                   |  0
 test/benchmarks/run_tests.sh                         |  1 +
 .../test => test/benchmarks}/test_result_analyzer.py |  2 +-
 5 files changed, 2 insertions(+), 18 deletions(-)
 delete mode 100755 benchmarks/test/run_tests.sh
 rename benchmarks/test/example.json => test/benchmarks/output-example.json (100%)
 rename {benchmarks/test => test/benchmarks}/test_result_analyzer.py (97%)

diff --git a/.circleci/common.sh b/.circleci/common.sh
index 642b596ceb4..28d8405c7e9 100755
--- a/.circleci/common.sh
+++ b/.circleci/common.sh
@@ -151,11 +151,6 @@ function run_torch_xla_python_tests() {
           # echo "Running MNIST Test"
           # python test/test_train_mp_mnist_amp.py --fake_data --num_epochs=1
         fi
-      elif [[ "$RUN_XLA_OP_TESTS1" == "xla_op1" ]]; then
-          # Benchmark tests.
-          # Only run on CPU, for xla_op1.
-          echo "Running Benchmark tests."
-          ./benchmarks/test/run_tests.sh
       fi
     fi
   popd
diff --git a/benchmarks/test/run_tests.sh b/benchmarks/test/run_tests.sh
deleted file mode 100755
index d3756734ecf..00000000000
--- a/benchmarks/test/run_tests.sh
+++ /dev/null
@@ -1,12 +0,0 @@
-BASEDIR="$(dirname $(dirname $(realpath $0)))"
-export PYTHONPATH="$BASEDIR"
-
-function run_test {
-    pushd "$BASEDIR"
-    python3 "$@"
-    popd
-}
-
-if [[ "$RUN_XLA_OP_TESTS1" == "xla_op1" ]]; then
-    run_test test/test_result_analyzer.py
-fi
diff --git a/benchmarks/test/example.json b/test/benchmarks/output-example.json
similarity index 100%
rename from benchmarks/test/example.json
rename to test/benchmarks/output-example.json
diff --git a/test/benchmarks/run_tests.sh b/test/benchmarks/run_tests.sh
index eef90b61347..7d404a7ee7f 100755
--- a/test/benchmarks/run_tests.sh
+++ b/test/benchmarks/run_tests.sh
@@ -42,6 +42,7 @@ function run_python_tests {
   python3 "$CDIR/test_experiment_runner.py"
   python3 "$CDIR/test_benchmark_experiment.py"
   python3 "$CDIR/test_benchmark_model.py"
+  python3 "$CDIR/test_result_analyzer.py"
 }
 
 function run_tests {
diff --git a/benchmarks/test/test_result_analyzer.py b/test/benchmarks/test_result_analyzer.py
similarity index 97%
rename from benchmarks/test/test_result_analyzer.py
rename to test/benchmarks/test_result_analyzer.py
index 7e1d8d2c9f0..d9237ee2e69 100644
--- a/benchmarks/test/test_result_analyzer.py
+++ b/test/benchmarks/test_result_analyzer.py
@@ -27,7 +27,7 @@ def get_dirname():
 @functools.cache
 def get_dataline():
   import json
-  example_json = os.path.join(get_dirname(), "example.json")
+  example_json = os.path.join(get_dirname(), "output-example.json")
   with open(example_json, "r") as f:
     return json.load(f)
 

From 9d8a436c3a466a1fe3a695b1590beca682278aa0 Mon Sep 17 00:00:00 2001
From: Yukio Siraichi <yukio.siraichi@gmail.com>
Date: Thu, 14 Dec 2023 08:44:29 -0300
Subject: [PATCH 2/2] Replace `functools.cache`.

---
 test/benchmarks/test_result_analyzer.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/test/benchmarks/test_result_analyzer.py b/test/benchmarks/test_result_analyzer.py
index d9237ee2e69..b636e143db1 100644
--- a/test/benchmarks/test_result_analyzer.py
+++ b/test/benchmarks/test_result_analyzer.py
@@ -1,5 +1,5 @@
 import argparse
-import functools
+import json
 import numpy
 import unittest
 import os
@@ -10,6 +10,8 @@
 fns_skip_head = (numpy.mean, numpy.std)
 fns_all = fns_whole + fns_skip_head
 
+_DATALINE = None
+
 
 def apply(fn, data):
   return fn(data)
@@ -19,17 +21,17 @@ def apply_skip_head(fn, data):
   return fn(data[1:])
 
 
-@functools.cache
 def get_dirname():
   return os.path.dirname(__file__)
 
 
-@functools.cache
 def get_dataline():
-  import json
-  example_json = os.path.join(get_dirname(), "output-example.json")
-  with open(example_json, "r") as f:
-    return json.load(f)
+  global _DATALINE
+  if _DATALINE is None:
+    example_json = os.path.join(get_dirname(), "output-example.json")
+    with open(example_json, "r") as f:
+      _DATALINE = json.load(f)
+  return _DATALINE
 
 
 class TestResultAnalyzer(unittest.TestCase):