[Clean up] Move utils/logs to r1/utils.

PiperOrigin-RevId: 309079916

official/r1/boosted_trees/train_higgs.py

@@ -45,16 +45,14 @@ from __future__ import print_function
 
 import os
 
-# pylint: disable=g-bad-import-order
-import numpy as np
 from absl import app as absl_app
 from absl import flags
+import numpy as np
 import tensorflow as tf
-# pylint: enable=g-bad-import-order
 
+from official.r1.utils.logs import logger
 from official.utils.flags import core as flags_core
 from official.utils.flags._conventions import help_wrap
-from official.utils.logs import logger
 
 NPZ_FILE = "HIGGS.csv.gz.npz"  # numpy compressed file containing "data" array
 

official/r1/mnist/mnist.py

@@ -24,8 +24,8 @@ from six.moves import range
 import tensorflow as tf
 
 from official.r1.mnist import dataset
+from official.r1.utils.logs import hooks_helper
 from official.utils.flags import core as flags_core
-from official.utils.logs import hooks_helper
 from official.utils.misc import distribution_utils
 from official.utils.misc import model_helpers
 

official/r1/ncf/ncf_estimator_main.py

@@ -17,7 +17,6 @@
 The NeuMF model assembles both MF and MLP models under the NCF framework. Check
 `neumf_model.py` for more details about the models.
 """
-
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -29,16 +28,17 @@ import math
 import multiprocessing
 import os
 import signal
-import typing
 
-# pylint: disable=g-bad-import-order
-import numpy as np
 from absl import app as absl_app
 from absl import flags
 from absl import logging
+import numpy as np
 import tensorflow as tf
-# pylint: enable=g-bad-import-order
+import typing
 
+from official.r1.utils.logs import hooks_helper
+from official.r1.utils.logs import logger
+from official.r1.utils.logs import mlperf_helper
 from official.recommendation import constants as rconst
 from official.recommendation import data_pipeline
 from official.recommendation import data_preprocessing
@@ -46,9 +46,6 @@ from official.recommendation import movielens
 from official.recommendation import ncf_common
 from official.recommendation import neumf_model
 from official.utils.flags import core as flags_core
-from official.utils.logs import hooks_helper
-from official.utils.logs import logger
-from official.utils.logs import mlperf_helper
 from official.utils.misc import distribution_utils
 from official.utils.misc import model_helpers
 

official/r1/resnet/cifar10_main.py

@@ -28,8 +28,8 @@ import tensorflow as tf
 
 from official.r1.resnet import resnet_model
 from official.r1.resnet import resnet_run_loop
+from official.r1.utils.logs import logger
 from official.utils.flags import core as flags_core
-from official.utils.logs import logger
 
 HEIGHT = 32
 WIDTH = 32

official/r1/resnet/estimator_benchmark.py

@@ -23,12 +23,12 @@ import time
 from absl import flags
 from absl import logging
 from absl.testing import flagsaver
-import tensorflow as tf  # pylint: disable=g-bad-import-order
+import tensorflow as tf
 
 from official.r1.resnet import cifar10_main as cifar_main
 from official.r1.resnet import imagenet_main
+from official.r1.utils.logs import hooks
 from official.utils.flags import core as flags_core
-from official.utils.logs import hooks
 
 IMAGENET_DATA_DIR_NAME = 'imagenet'
 CIFAR_DATA_DIR_NAME = 'cifar-10-batches-bin'

official/r1/resnet/imagenet_main.py

@@ -29,8 +29,8 @@ import tensorflow as tf
 from official.r1.resnet import imagenet_preprocessing
 from official.r1.resnet import resnet_model
 from official.r1.resnet import resnet_run_loop
+from official.r1.utils.logs import logger
 from official.utils.flags import core as flags_core
-from official.utils.logs import logger
 
 DEFAULT_IMAGE_SIZE = 224
 NUM_CHANNELS = 3

official/r1/resnet/resnet_run_loop.py

@@ -35,9 +35,9 @@ import tensorflow as tf
 from official.r1.resnet import imagenet_preprocessing
 from official.r1.resnet import resnet_model
 from official.r1.utils import export
+from official.r1.utils.logs import hooks_helper
+from official.r1.utils.logs import logger
 from official.utils.flags import core as flags_core
-from official.utils.logs import hooks_helper
-from official.utils.logs import logger
 from official.utils.misc import distribution_utils
 from official.utils.misc import model_helpers
 

official/r1/transformer/transformer_main.py

@@ -43,8 +43,8 @@ from official.nlp.transformer import compute_bleu
 from official.nlp.transformer.utils import metrics
 from official.nlp.transformer.utils import tokenizer
 from official.utils.flags import core as flags_core
-from official.utils.logs import hooks_helper
-from official.utils.logs import logger
+from official.r1.utils.logs import hooks_helper
+from official.r1.utils.logs import logger
 from official.utils.misc import distribution_utils
 from official.utils.misc import model_helpers
 

official/utils/logs/cloud_lib_test.py → official/r1/utils/logs/cloud_lib_test.py

@@ -24,7 +24,7 @@ import unittest
 import mock
 import requests
 
-from official.utils.logs import cloud_lib
+from official.r1.utils.logs import cloud_lib
 
 
 class CloudLibTest(unittest.TestCase):

official/utils/logs/hooks.py → official/r1/utils/logs/hooks.py

@@ -22,7 +22,7 @@ from __future__ import print_function
 
 import tensorflow as tf  # pylint: disable=g-bad-import-order
 
-from official.utils.logs import logger
+from official.r1.utils.logs import logger
 
 
 class ExamplesPerSecondHook(tf.estimator.SessionRunHook):

official/utils/logs/hooks_helper.py → official/r1/utils/logs/hooks_helper.py

@@ -27,9 +27,9 @@ from __future__ import print_function
 import tensorflow as tf  # pylint: disable=g-bad-import-order
 from absl import logging
 
-from official.utils.logs import hooks
-from official.utils.logs import logger
-from official.utils.logs import metric_hook
+from official.r1.utils.logs import hooks
+from official.r1.utils.logs import logger
+from official.r1.utils.logs import metric_hook
 
 _TENSORS_TO_LOG = dict((x, x) for x in ['learning_rate',
                                         'cross_entropy',

official/utils/logs/hooks_test.py → official/r1/utils/logs/hooks_test.py

@@ -24,8 +24,8 @@ import time
 from absl import logging
 import tensorflow as tf  # pylint: disable=g-bad-import-order
 
-from official.utils.logs import hooks
-from official.utils.testing import mock_lib
+from official.r1.utils.logs import hooks
+from official.r1.utils.logs import mock_lib
 
 logging.set_verbosity(logging.DEBUG)
 

official/utils/logs/logger.py → official/r1/utils/logs/logger.py

@@ -31,13 +31,12 @@ import os
 import threading
 import uuid
 
-from six.moves import _thread as thread
 from absl import flags
+from absl import logging
+from six.moves import _thread as thread
 import tensorflow as tf
 from tensorflow.python.client import device_lib
-from absl import logging
-
-from official.utils.logs import cloud_lib
+from official.r1.utils.logs import cloud_lib
 
 METRIC_LOG_FILE_NAME = "metric.log"
 BENCHMARK_RUN_LOG_FILE_NAME = "benchmark_run.log"

official/utils/logs/logger_test.py → official/r1/utils/logs/logger_test.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 """Tests for benchmark logger."""
-
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -25,19 +24,13 @@ import tempfile
 import time
 import unittest
 
-import mock
-from absl.testing import flagsaver
-import tensorflow as tf  # pylint: disable=g-bad-import-order
 from absl import logging
+from absl.testing import flagsaver
+import tensorflow as tf
 
-try:
-  from google.cloud import bigquery
-except ImportError:
-  bigquery = None
-
-from official.utils.misc import keras_utils
+from official.r1.utils.logs import logger
 from official.utils.flags import core as flags_core
-from official.utils.logs import logger
+from official.utils.misc import keras_utils
 
 
 class BenchmarkLoggerTest(tf.test.TestCase):
@@ -67,23 +60,6 @@ class BenchmarkLoggerTest(tf.test.TestCase):
         self.assertIsInstance(logger.get_benchmark_logger(),
                               logger.BenchmarkFileLogger)
 
-  @mock.patch("official.utils.logs.logger.config_benchmark_logger")
-  def test_benchmark_context(self, mock_config_benchmark_logger):
-    mock_logger = mock.MagicMock()
-    mock_config_benchmark_logger.return_value = mock_logger
-    with logger.benchmark_context(None):
-      logging.info("start benchmarking")
-    mock_logger.on_finish.assert_called_once_with(logger.RUN_STATUS_SUCCESS)
-
-  @mock.patch("official.utils.logs.logger.config_benchmark_logger")
-  def test_benchmark_context_failure(self, mock_config_benchmark_logger):
-    mock_logger = mock.MagicMock()
-    mock_config_benchmark_logger.return_value = mock_logger
-    with self.assertRaises(RuntimeError):
-      with logger.benchmark_context(None):
-        raise RuntimeError("training error")
-    mock_logger.on_finish.assert_called_once_with(logger.RUN_STATUS_FAILURE)
-
 
 class BaseBenchmarkLoggerTest(tf.test.TestCase):
 
@@ -212,24 +188,6 @@ class BenchmarkFileLoggerTest(tf.test.TestCase):
     metric_log = os.path.join(log_dir, "metric.log")
     self.assertFalse(tf.io.gfile.exists(metric_log))
 
-  @mock.patch("official.utils.logs.logger._gather_run_info")
-  def test_log_run_info(self, mock_gather_run_info):
-    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
-    log = logger.BenchmarkFileLogger(log_dir)
-    run_info = {"model_name": "model_name",
-                "dataset": "dataset_name",
-                "run_info": "run_value"}
-    mock_gather_run_info.return_value = run_info
-    log.log_run_info("model_name", "dataset_name", {})
-
-    run_log = os.path.join(log_dir, "benchmark_run.log")
-    self.assertTrue(tf.io.gfile.exists(run_log))
-    with tf.io.gfile.GFile(run_log) as f:
-      run_info = json.loads(f.readline())
-      self.assertEqual(run_info["model_name"], "model_name")
-      self.assertEqual(run_info["dataset"], "dataset_name")
-      self.assertEqual(run_info["run_info"], "run_value")
-
   def test_collect_tensorflow_info(self):
     run_info = {}
     logger._collect_tensorflow_info(run_info)

official/utils/logs/metric_hook_test.py → official/r1/utils/logs/metric_hook_test.py

@@ -24,8 +24,8 @@ import time
 import tensorflow as tf  # pylint: disable=g-bad-import-order
 from tensorflow.python.training import monitored_session  # pylint: disable=g-bad-import-order
 
-from official.utils.logs import metric_hook
-from official.utils.testing import mock_lib
+from official.r1.utils.logs import metric_hook
+from official.r1.utils.logs import mock_lib
 
 
 class LoggingMetricHookTest(tf.test.TestCase):

official/utils/logs/mlperf_helper.py → official/r1/utils/logs/mlperf_helper.py

@@ -30,8 +30,8 @@ import os
 import re
 import subprocess
 import sys
-import typing
 from absl import logging
+import typing
 # pylint:disable=logging-format-interpolation