logger_test.py 14.3 KB
Newer Older
Scott Zhu's avatar
Scott Zhu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Tests for benchmark logger."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import json
import os
import tempfile
25
import time
26
import unittest
Scott Zhu's avatar
Scott Zhu committed
27

28
29
import mock
from absl.testing import flagsaver
Karmel Allison's avatar
Karmel Allison committed
30
import tensorflow as tf  # pylint: disable=g-bad-import-order
31
from absl import logging
Scott Zhu's avatar
Scott Zhu committed
32

33
34
35
36
37
try:
  from google.cloud import bigquery
except ImportError:
  bigquery = None

38
from official.utils.misc import keras_utils
39
from official.utils.flags import core as flags_core
40
from official.utils.logs import logger
Scott Zhu's avatar
Scott Zhu committed
41
42
43
44


class BenchmarkLoggerTest(tf.test.TestCase):

45
46
47
48
49
  @classmethod
  def setUpClass(cls):  # pylint: disable=invalid-name
    super(BenchmarkLoggerTest, cls).setUpClass()
    flags_core.define_benchmark()

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
50
  def test_get_default_benchmark_logger(self):
51
    with flagsaver.flagsaver(benchmark_logger_type="foo"):
52
53
      self.assertIsInstance(logger.get_benchmark_logger(),
                            logger.BaseBenchmarkLogger)
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
54
55

  def test_config_base_benchmark_logger(self):
56
    with flagsaver.flagsaver(benchmark_logger_type="BaseBenchmarkLogger"):
57
58
59
      logger.config_benchmark_logger()
      self.assertIsInstance(logger.get_benchmark_logger(),
                            logger.BaseBenchmarkLogger)
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
60
61

  def test_config_benchmark_file_logger(self):
62
63
    # Set the benchmark_log_dir first since the benchmark_logger_type will need
    # the value to be set when it does the validation.
64
65
    with flagsaver.flagsaver(benchmark_log_dir="/tmp"):
      with flagsaver.flagsaver(benchmark_logger_type="BenchmarkFileLogger"):
66
67
68
69
        logger.config_benchmark_logger()
        self.assertIsInstance(logger.get_benchmark_logger(),
                              logger.BenchmarkFileLogger)

70
  @unittest.skipIf(bigquery is None, "Bigquery dependency is not installed.")
71
72
  @mock.patch.object(bigquery, "Client")
  def test_config_benchmark_bigquery_logger(self, mock_bigquery_client):
73
    with flagsaver.flagsaver(benchmark_logger_type="BenchmarkBigQueryLogger"):
74
75
76
      logger.config_benchmark_logger()
      self.assertIsInstance(logger.get_benchmark_logger(),
                            logger.BenchmarkBigQueryLogger)
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
77

78
79
80
81
82
  @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):
83
      logging.info("start benchmarking")
84
85
86
87
88
89
90
91
92
93
94
    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)

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
95
96
97
98
99

class BaseBenchmarkLoggerTest(tf.test.TestCase):

  def setUp(self):
    super(BaseBenchmarkLoggerTest, self).setUp()
100
    self._actual_log = logging.info
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
101
102
103
104
105
106
    self.logged_message = None

    def mock_log(*args, **kwargs):
      self.logged_message = args
      self._actual_log(*args, **kwargs)

107
    logging.info = mock_log
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
108
109
110

  def tearDown(self):
    super(BaseBenchmarkLoggerTest, self).tearDown()
111
    logging.info = self._actual_log
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
112
113
114
115
116
117
118
119
120
121
122

  def test_log_metric(self):
    log = logger.BaseBenchmarkLogger()
    log.log_metric("accuracy", 0.999, global_step=1e4, extras={"name": "value"})

    expected_log_prefix = "Benchmark metric:"
    self.assertRegexpMatches(str(self.logged_message), expected_log_prefix)


class BenchmarkFileLoggerTest(tf.test.TestCase):

123
  def setUp(self):
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
124
    super(BenchmarkFileLoggerTest, self).setUp()
125
126
127
128
129
130
    # Avoid pulling extra env vars from test environment which affects the test
    # result, eg. Kokoro test has a TF_PKG env which affect the test case
    # test_collect_tensorflow_environment_variables()
    self.original_environ = dict(os.environ)
    os.environ.clear()

Scott Zhu's avatar
Scott Zhu committed
131
  def tearDown(self):
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
132
    super(BenchmarkFileLoggerTest, self).tearDown()
133
    tf.io.gfile.rmtree(self.get_temp_dir())
134
135
    os.environ.clear()
    os.environ.update(self.original_environ)
Scott Zhu's avatar
Scott Zhu committed
136
137
138

  def test_create_logging_dir(self):
    non_exist_temp_dir = os.path.join(self.get_temp_dir(), "unknown_dir")
139
    self.assertFalse(tf.io.gfile.isdir(non_exist_temp_dir))
Scott Zhu's avatar
Scott Zhu committed
140

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
141
    logger.BenchmarkFileLogger(non_exist_temp_dir)
142
    self.assertTrue(tf.io.gfile.isdir(non_exist_temp_dir))
Scott Zhu's avatar
Scott Zhu committed
143
144
145

  def test_log_metric(self):
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
146
    log = logger.BenchmarkFileLogger(log_dir)
Scott Zhu's avatar
Scott Zhu committed
147
148
149
    log.log_metric("accuracy", 0.999, global_step=1e4, extras={"name": "value"})

    metric_log = os.path.join(log_dir, "metric.log")
150
151
    self.assertTrue(tf.io.gfile.exists(metric_log))
    with tf.io.gfile.GFile(metric_log) as f:
Scott Zhu's avatar
Scott Zhu committed
152
153
154
155
156
      metric = json.loads(f.readline())
      self.assertEqual(metric["name"], "accuracy")
      self.assertEqual(metric["value"], 0.999)
      self.assertEqual(metric["unit"], None)
      self.assertEqual(metric["global_step"], 1e4)
157
      self.assertEqual(metric["extras"], [{"name": "name", "value": "value"}])
Scott Zhu's avatar
Scott Zhu committed
158
159
160

  def test_log_multiple_metrics(self):
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
161
    log = logger.BenchmarkFileLogger(log_dir)
Scott Zhu's avatar
Scott Zhu committed
162
163
164
165
    log.log_metric("accuracy", 0.999, global_step=1e4, extras={"name": "value"})
    log.log_metric("loss", 0.02, global_step=1e4)

    metric_log = os.path.join(log_dir, "metric.log")
166
167
    self.assertTrue(tf.io.gfile.exists(metric_log))
    with tf.io.gfile.GFile(metric_log) as f:
Scott Zhu's avatar
Scott Zhu committed
168
169
170
171
172
      accuracy = json.loads(f.readline())
      self.assertEqual(accuracy["name"], "accuracy")
      self.assertEqual(accuracy["value"], 0.999)
      self.assertEqual(accuracy["unit"], None)
      self.assertEqual(accuracy["global_step"], 1e4)
173
      self.assertEqual(accuracy["extras"], [{"name": "name", "value": "value"}])
Scott Zhu's avatar
Scott Zhu committed
174
175
176
177
178
179

      loss = json.loads(f.readline())
      self.assertEqual(loss["name"], "loss")
      self.assertEqual(loss["value"], 0.02)
      self.assertEqual(loss["unit"], None)
      self.assertEqual(loss["global_step"], 1e4)
180
      self.assertEqual(loss["extras"], [])
Scott Zhu's avatar
Scott Zhu committed
181

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
182
  def test_log_non_number_value(self):
Scott Zhu's avatar
Scott Zhu committed
183
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
184
    log = logger.BenchmarkFileLogger(log_dir)
Scott Zhu's avatar
Scott Zhu committed
185
186
187
188
    const = tf.constant(1)
    log.log_metric("accuracy", const)

    metric_log = os.path.join(log_dir, "metric.log")
189
    self.assertFalse(tf.io.gfile.exists(metric_log))
Scott Zhu's avatar
Scott Zhu committed
190

191
  def test_log_evaluation_result(self):
192
193
194
    eval_result = {"loss": 0.46237424,
                   "global_step": 207082,
                   "accuracy": 0.9285}
195
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
196
197
    log = logger.BenchmarkFileLogger(log_dir)
    log.log_evaluation_result(eval_result)
198
199

    metric_log = os.path.join(log_dir, "metric.log")
200
201
    self.assertTrue(tf.io.gfile.exists(metric_log))
    with tf.io.gfile.GFile(metric_log) as f:
202
203
204
205
206
207
      accuracy = json.loads(f.readline())
      self.assertEqual(accuracy["name"], "accuracy")
      self.assertEqual(accuracy["value"], 0.9285)
      self.assertEqual(accuracy["unit"], None)
      self.assertEqual(accuracy["global_step"], 207082)

208
209
210
211
212
213
      loss = json.loads(f.readline())
      self.assertEqual(loss["name"], "loss")
      self.assertEqual(loss["value"], 0.46237424)
      self.assertEqual(loss["unit"], None)
      self.assertEqual(loss["global_step"], 207082)

214
215
216
  def test_log_evaluation_result_with_invalid_type(self):
    eval_result = "{'loss': 0.46237424, 'global_step': 207082}"
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
217
218
    log = logger.BenchmarkFileLogger(log_dir)
    log.log_evaluation_result(eval_result)
219
220

    metric_log = os.path.join(log_dir, "metric.log")
221
    self.assertFalse(tf.io.gfile.exists(metric_log))
222

223
224
225
226
227
228
229
230
231
232
233
  @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")
234
235
    self.assertTrue(tf.io.gfile.exists(run_log))
    with tf.io.gfile.GFile(run_log) as f:
236
237
238
239
240
      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")

241
242
243
244
  def test_collect_tensorflow_info(self):
    run_info = {}
    logger._collect_tensorflow_info(run_info)
    self.assertNotEqual(run_info["tensorflow_version"], {})
245
246
247
248
    self.assertEqual(run_info["tensorflow_version"]["version"],
                     tf.version.VERSION)
    self.assertEqual(run_info["tensorflow_version"]["git_hash"],
                     tf.version.GIT_VERSION)
249

250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
  def test_collect_run_params(self):
    run_info = {}
    run_parameters = {
        "batch_size": 32,
        "synthetic_data": True,
        "train_epochs": 100.00,
        "dtype": "fp16",
        "resnet_size": 50,
        "random_tensor": tf.constant(2.0)
    }
    logger._collect_run_params(run_info, run_parameters)
    self.assertEqual(len(run_info["run_parameters"]), 6)
    self.assertEqual(run_info["run_parameters"][0],
                     {"name": "batch_size", "long_value": 32})
    self.assertEqual(run_info["run_parameters"][1],
                     {"name": "dtype", "string_value": "fp16"})
266
267
268
269
270
271
    v1_tensor = {"name": "random_tensor", "string_value":
                     "Tensor(\"Const:0\", shape=(), dtype=float32)"}
    v2_tensor = {"name": "random_tensor", "string_value":
                     "tf.Tensor(2.0, shape=(), dtype=float32)"}
    self.assertIn(run_info["run_parameters"][2], [v1_tensor, v2_tensor])

272

273
274
275
276
277
278
279
    self.assertEqual(run_info["run_parameters"][3],
                     {"name": "resnet_size", "long_value": 50})
    self.assertEqual(run_info["run_parameters"][4],
                     {"name": "synthetic_data", "bool_value": "True"})
    self.assertEqual(run_info["run_parameters"][5],
                     {"name": "train_epochs", "float_value": 100.00})

280
281
  def test_collect_tensorflow_environment_variables(self):
    os.environ["TF_ENABLE_WINOGRAD_NONFUSED"] = "1"
282
283
    os.environ["TF_OTHER"] = "2"
    os.environ["OTHER"] = "3"
284
285
286
287

    run_info = {}
    logger._collect_tensorflow_environment_variables(run_info)
    self.assertIsNotNone(run_info["tensorflow_environment_variables"])
288
289
290
291
292
293
    expected_tf_envs = [
        {"name": "TF_ENABLE_WINOGRAD_NONFUSED", "value": "1"},
        {"name": "TF_OTHER", "value": "2"},
    ]
    self.assertEqual(run_info["tensorflow_environment_variables"],
                     expected_tf_envs)
294
295

  def test_collect_memory_info(self):
296
    run_info = {"machine_config": {}}
297
    logger._collect_memory_info(run_info)
298
299
    self.assertIsNotNone(run_info["machine_config"]["memory_total"])
    self.assertIsNotNone(run_info["machine_config"]["memory_available"])
300

301

302
@unittest.skipIf(bigquery is None, "Bigquery dependency is not installed.")
303
304
305
306
307
308
309
310
311
312
313
314
class BenchmarkBigQueryLoggerTest(tf.test.TestCase):

  def setUp(self):
    super(BenchmarkBigQueryLoggerTest, self).setUp()
    # Avoid pulling extra env vars from test environment which affects the test
    # result, eg. Kokoro test has a TF_PKG env which affect the test case
    # test_collect_tensorflow_environment_variables()
    self.original_environ = dict(os.environ)
    os.environ.clear()

    self.mock_bq_uploader = mock.MagicMock()
    self.logger = logger.BenchmarkBigQueryLogger(
315
316
        self.mock_bq_uploader, "dataset", "run_table", "run_status_table",
        "metric_table", "run_id")
317
318
319

  def tearDown(self):
    super(BenchmarkBigQueryLoggerTest, self).tearDown()
320
    tf.io.gfile.rmtree(self.get_temp_dir())
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
    os.environ.clear()
    os.environ.update(self.original_environ)

  def test_log_metric(self):
    self.logger.log_metric(
        "accuracy", 0.999, global_step=1e4, extras={"name": "value"})
    expected_metric_json = [{
        "name": "accuracy",
        "value": 0.999,
        "unit": None,
        "global_step": 1e4,
        "timestamp": mock.ANY,
        "extras": [{"name": "name", "value": "value"}]
    }]
    # log_metric will call upload_benchmark_metric_json in a separate thread.
    # Give it some grace period for the new thread before assert.
    time.sleep(1)
    self.mock_bq_uploader.upload_benchmark_metric_json.assert_called_once_with(
        "dataset", "metric_table", "run_id", expected_metric_json)

341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
  @mock.patch("official.utils.logs.logger._gather_run_info")
  def test_log_run_info(self, mock_gather_run_info):
    run_info = {"model_name": "model_name",
                "dataset": "dataset_name",
                "run_info": "run_value"}
    mock_gather_run_info.return_value = run_info
    self.logger.log_run_info("model_name", "dataset_name", {})
    # log_metric will call upload_benchmark_metric_json in a separate thread.
    # Give it some grace period for the new thread before assert.
    time.sleep(1)
    self.mock_bq_uploader.upload_benchmark_run_json.assert_called_once_with(
        "dataset", "run_table", "run_id", run_info)
    self.mock_bq_uploader.insert_run_status.assert_called_once_with(
        "dataset", "run_status_table", "run_id", "running")

  def test_on_finish(self):
    self.logger.on_finish(logger.RUN_STATUS_SUCCESS)
    # log_metric will call upload_benchmark_metric_json in a separate thread.
    # Give it some grace period for the new thread before assert.
    time.sleep(1)
    self.mock_bq_uploader.update_run_status.assert_called_once_with(
        "dataset", "run_status_table", "run_id", logger.RUN_STATUS_SUCCESS)

364

Scott Zhu's avatar
Scott Zhu committed
365
366
if __name__ == "__main__":
  tf.test.main()