From 1fbe9351bf8aa4e2099c10a8ecab969db7cde4c5 Mon Sep 17 00:00:00 2001
From: Yukio Siraichi <yukio.siraichi@gmail.com>
Date: Thu, 14 Dec 2023 13:21:50 -0300
Subject: [PATCH] Move `test_result_analyzer.py` to the test directory. (#6130)

---
 .circleci/common.sh                              |  5 -----
 benchmarks/test/run_tests.sh                     | 12 ------------
 .../benchmarks/output-example.json               |  0
 test/benchmarks/run_tests.sh                     |  1 +
 .../benchmarks}/test_result_analyzer.py          | 16 +++++++++-------
 5 files changed, 10 insertions(+), 24 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 (92%)

diff --git a/.circleci/common.sh b/.circleci/common.sh
index 642b596ceb4b..28d8405c7e99 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 d3756734ecf8..000000000000
--- 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 eef90b613478..7d404a7ee7ff 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 92%
rename from benchmarks/test/test_result_analyzer.py
rename to test/benchmarks/test_result_analyzer.py
index 7e1d8d2c9f0e..b636e143db19 100644
--- a/benchmarks/test/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(), "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):