chore: import upstream snapshot with attribution

tensorflow/python/profiler/integration_test/BUILD

@@ -0,0 +1,39 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "mnist_testing_utils",
+    srcs = ["mnist_testing_utils.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python:extra_py_tests_deps",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "profiler_api_test",
+    srcs = ["profiler_api_test.py"],
+    tags = [
+        "no_oss",  # TODO(b/283175845)
+        "no_pip",
+        "no_windows",  # TODO(b/192257727)
+    ],
+    deps = [
+        ":mnist_testing_utils",
+        "//tensorflow/python/distribute:collective_all_reduce_strategy",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/profiler:profiler_client",
+        "//tensorflow/python/profiler:profiler_v2",
+        "@pypi//portpicker",
+    ],
+)

tensorflow/python/profiler/integration_test/mnist_testing_utils.py

@@ -0,0 +1,68 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""A simple MNIST model for testing multi-worker distribution strategies with Keras."""
+
+import tensorflow as tf
+
+
+def mnist_synthetic_dataset(batch_size, steps_per_epoch):
+  """Generate synthetic MNIST dataset for testing."""
+  # train dataset
+  x_train = tf.ones([batch_size * steps_per_epoch, 28, 28, 1],
+                    dtype=tf.dtypes.float32)
+  y_train = tf.ones([batch_size * steps_per_epoch, 1], dtype=tf.dtypes.int32)
+  train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+  train_ds = train_ds.repeat()
+  # train_ds = train_ds.shuffle(100)
+  train_ds = train_ds.batch(64, drop_remainder=True)
+
+  # eval dataset
+  x_test = tf.random.uniform([10000, 28, 28, 1], dtype=tf.dtypes.float32)
+  y_test = tf.random.uniform([10000, 1],
+                             minval=0,
+                             maxval=9,
+                             dtype=tf.dtypes.int32)
+  eval_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+  eval_ds = eval_ds.batch(64, drop_remainder=True)
+
+  return train_ds, eval_ds
+
+
+def get_mnist_model(input_shape):
+  """Define a deterministically-initialized CNN model for MNIST testing."""
+  inputs = tf.keras.Input(shape=input_shape)
+  x = tf.keras.layers.Conv2D(
+      32,
+      kernel_size=(3, 3),
+      activation="relu",
+      kernel_initializer=tf.keras.initializers.TruncatedNormal(seed=99))(
+          inputs)
+  x = tf.keras.layers.BatchNormalization()(x)
+  x = tf.keras.layers.Flatten()(x) + tf.keras.layers.Flatten()(x)
+  x = tf.keras.layers.Dense(
+      10,
+      activation="softmax",
+      kernel_initializer=tf.keras.initializers.TruncatedNormal(seed=99))(
+          x)
+  model = tf.keras.Model(inputs=inputs, outputs=x)
+
+  # TODO(yuefengz): optimizer with slot variables doesn't work because of
+  # optimizer's bug.
+  # TODO(yuefengz): we should not allow non-v2 optimizer.
+  model.compile(
+      loss=tf.keras.losses.sparse_categorical_crossentropy,
+      optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
+      metrics=["accuracy"])
+  return model

tensorflow/python/profiler/integration_test/profiler_api_test.py

@@ -0,0 +1,148 @@
+# Copyright 2020 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 tf 2.x profiler."""
+
+import glob
+import os
+import threading
+
+import portpicker
+
+from tensorflow.python.distribute import collective_all_reduce_strategy as collective_strategy
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.eager import context
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.profiler import profiler_client
+from tensorflow.python.profiler import profiler_v2 as profiler
+from tensorflow.python.profiler.integration_test import mnist_testing_utils
+
+
+def _model_setup():
+  """Set up a MNIST Keras model for testing purposes.
+
+  Builds a MNIST Keras model and returns model information.
+
+  Returns:
+    A tuple of (batch_size, steps, train_dataset, mode)
+  """
+  context.set_log_device_placement(True)
+  batch_size = 64
+  steps = 2
+  with collective_strategy.CollectiveAllReduceStrategy().scope():
+    # TODO(b/142509827): In rare cases this errors out at C++ level with the
+    # "Connect failed" error message.
+    train_ds, _ = mnist_testing_utils.mnist_synthetic_dataset(batch_size, steps)
+    model = mnist_testing_utils.get_mnist_model((28, 28, 1))
+  return batch_size, steps, train_ds, model
+
+
+def _make_temp_log_dir(test_obj):
+  return test_obj.get_temp_dir()
+
+
+class ProfilerApiTest(test_util.TensorFlowTestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.worker_start = threading.Event()
+    self.profile_done = False
+
+  def _check_xspace_pb_exist(self, logdir):
+    path = os.path.join(logdir, 'plugins', 'profile', '*', '*.xplane.pb')
+    self.assertEqual(1, len(glob.glob(path)),
+                     'Expected one path match: ' + path)
+
+  def test_single_worker_no_profiling(self):
+    """Test single worker without profiling."""
+
+    _, steps, train_ds, model = _model_setup()
+
+    model.fit(x=train_ds, epochs=2, steps_per_epoch=steps)
+
+  def test_single_worker_sampling_mode(self, delay_ms=None):
+    """Test single worker sampling mode."""
+
+    def on_worker(port, worker_start):
+      logging.info('worker starting server on {}'.format(port))
+      profiler.start_server(port)
+      _, steps, train_ds, model = _model_setup()
+      worker_start.set()
+      while True:
+        model.fit(x=train_ds, epochs=2, steps_per_epoch=steps)
+        if self.profile_done:
+          break
+
+    def on_profile(port, logdir, worker_start):
+      # Request for 30 milliseconds of profile.
+      duration_ms = 30
+
+      worker_start.wait()
+      options = profiler.ProfilerOptions(
+          host_tracer_level=2,
+          python_tracer_level=0,
+          device_tracer_level=1,
+          delay_ms=delay_ms,
+      )
+
+      profiler_client.trace('localhost:{}'.format(port), logdir, duration_ms,
+                            '', 100, options)
+
+      self.profile_done = True
+
+    logdir = self.get_temp_dir()
+    port = portpicker.pick_unused_port()
+    thread_profiler = threading.Thread(
+        target=on_profile, args=(port, logdir, self.worker_start))
+    thread_worker = threading.Thread(
+        target=on_worker, args=(port, self.worker_start))
+    thread_worker.start()
+    thread_profiler.start()
+    thread_profiler.join()
+    thread_worker.join(120)
+    self._check_xspace_pb_exist(logdir)
+
+  def test_single_worker_sampling_mode_short_delay(self):
+    """Test single worker sampling mode with a short delay.
+
+    Expect that requested delayed start time will arrive late, and a subsequent
+    retry will issue an immediate start.
+    """
+
+    self.test_single_worker_sampling_mode(delay_ms=1)
+
+  def test_single_worker_sampling_mode_long_delay(self):
+    """Test single worker sampling mode with a long delay."""
+
+    self.test_single_worker_sampling_mode(delay_ms=1000)
+
+  def test_single_worker_programmatic_mode(self):
+    """Test single worker programmatic mode."""
+    logdir = self.get_temp_dir()
+
+    options = profiler.ProfilerOptions(
+        host_tracer_level=2,
+        python_tracer_level=0,
+        device_tracer_level=1,
+    )
+    profiler.start(logdir, options)
+    _, steps, train_ds, model = _model_setup()
+    model.fit(x=train_ds, epochs=2, steps_per_epoch=steps)
+    profiler.stop()
+    self._check_xspace_pb_exist(logdir)
+
+
+if __name__ == '__main__':
+  multi_process_runner.test_main()