Update references of tensorflow/python:platform to use the public API if available and remove references to the deprecated tensorflow/python:platform target

PiperOrigin-RevId: 524913413

Author: BrianWieder

tensorflow_estimator/python/estimator/BUILD

@@ -768,6 +768,7 @@ py_library(
     name = "analytics_tools",
     srcs = ["tools/analytics.py"],
     srcs_version = "PY3",
+    deps = ["//tensorflow_estimator/python/estimator:expect_tensorflow_installed"],
 )
 
 py_test(
@@ -1906,7 +1907,7 @@ py_test(
     srcs_version = "PY3",
     deps = [
         ":estimator_export",
-        ":expect_tensorflow_installed",
+        "//tensorflow_estimator/python/estimator:expect_tensorflow_installed",
     ],
 )
 

tensorflow_estimator/python/estimator/early_stopping.py

@@ -19,7 +19,6 @@
 import os
 
 import tensorflow as tf
-from tensorflow.python.platform import tf_logging
 from tensorflow_estimator.python.estimator import estimator as estimator_lib
 from tensorflow_estimator.python.estimator.estimator_export import estimator_export
 
@@ -578,15 +577,15 @@ def after_run(self, run_context, run_values):
     should_early_stop = run_values.results['stop_var']
 
     if should_early_stop > 0:
-      tf_logging.info('Early stopping requested, suspending run.')
+      tf.compat.v1.logging.info('Early stopping requested, suspending run.')
       run_context.request_stop()
       return
     if self._timer.should_trigger_for_step(global_step):
       self._timer.update_last_triggered_step(global_step)
       if self._should_stop_fn():
         run_context.session.run(
             self._stop_op, feed_dict={self._stop_placeholder: 1})
-        tf_logging.info('Requesting early stopping at global step %d',
+        tf.compat.v1.logging.info('Requesting early stopping at global step %d',
                         global_step)
       else:
         run_context.session.run(

tensorflow_estimator/python/estimator/estimator.py

@@ -35,7 +35,6 @@
 from tensorflow.python.eager import context
 from tensorflow.python.eager import monitoring
 from tensorflow.python.framework import ops
-from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.profiler import trace
 from tensorflow.python.saved_model import path_helpers
 from tensorflow.python.summary import summary
@@ -349,15 +348,17 @@ def train(self,
       if max_steps is not None:
         start_step = _load_global_step_from_checkpoint_dir(self._model_dir)
         if max_steps <= start_step:
-          logging.info('Skipping training since max_steps has already saved.')
+          tf.compat.v1.logging.info(
+              'Skipping training since max_steps has already saved.'
+          )
           return self
 
       hooks = _check_hooks_type(hooks)
       hooks.extend(self._convert_train_steps_to_hooks(steps, max_steps))
 
       saving_listeners = _check_listeners_type(saving_listeners)
       loss = self._train_model(input_fn, hooks, saving_listeners)
-      logging.info('Loss for final step: %s.', loss)
+      tf.compat.v1.logging.info('Loss for final step: %s.', loss)
       return self
 
   def _convert_train_steps_to_hooks(self, steps, max_steps):
@@ -1021,9 +1022,11 @@ def _get_features_from_input_fn(self, input_fn, mode):
 
   def _validate_features_in_predict_input(self, result):
     if not _has_dataset_or_queue_runner(result):
-      logging.warning('Input graph does not use tf.data.Dataset or contain a '
-                      'QueueRunner. That means predict yields forever. '
-                      'This is probably a mistake.')
+      tf.compat.v1.logging.warning(
+          'Input graph does not use tf.data.Dataset or contain a '
+          'QueueRunner. That means predict yields forever. '
+          'This is probably a mistake.'
+      )
 
   def _get_iterator_from_input_fn(self, input_fn, mode, distribution=None):
     """Calls `input_fn` and returns an iterator."""
@@ -1169,9 +1172,9 @@ def _call_model_fn(self, features, labels, mode, config):
     if 'config' in model_fn_args:
       kwargs['config'] = config
 
-    logging.info('Calling model_fn.')
+    tf.compat.v1.logging.info('Calling model_fn.')
     model_fn_results = self._model_fn(features=features, **kwargs)
-    logging.info('Done calling model_fn.')
+    tf.compat.v1.logging.info('Done calling model_fn.')
 
     if not isinstance(model_fn_results, model_fn_lib.EstimatorSpec):
       raise ValueError('model_fn should return an EstimatorSpec.')

tensorflow_estimator/python/estimator/estimator_export_test.py

@@ -15,8 +15,8 @@
 """estimator_export tests."""
 
 import sys
+import tensorflow as tf
 
-from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.util import tf_export
 # pylint: disable=g-deprecated-tf-checker
@@ -27,7 +27,7 @@ class TestClass(object):
   pass
 
 
-class ValidateExportTest(test.TestCase):
+class ValidateExportTest(tf.test.TestCase):
   """Tests for estimator_export class."""
 
   def setUp(self):
@@ -49,7 +49,9 @@ def testExportClassInEstimator(self):
     self.assertNotIn('_tf_api_names', TestClass.__dict__)
     self.assertEqual(['estimator.TestClass'], tf_export.get_v1_names(TestClass))
 
-  @test.mock.patch.object(logging, 'warning', autospec=True)
+  @tf.compat.v1.test.mock.patch.object(
+      logging, 'warning', autospec=True
+  )
   def testExportDeprecated(self, mock_warning):
     export_decorator = estimator_export.estimator_export('estimator.TestClass')
     export_decorator(TestClass)
@@ -63,4 +65,4 @@ def testExportDeprecated(self, mock_warning):
 
 
 if __name__ == '__main__':
-  test.main()
+  tf.test.main()

tensorflow_estimator/python/estimator/estimator_test.py

@@ -1925,7 +1925,7 @@ def _model_fn(features, labels, mode):
 
     est = estimator.EstimatorV2(model_fn=_model_fn)
     est.train(dummy_input_fn, steps=1)
-    with tf.test.mock.patch.object(logging, 'warning') as mock_log:
+    with tf.test.mock.patch.object(tf.logging, 'warning') as mock_log:
       next(est.predict(dummy_input_fn))
       self.assertRegexpMatches(
           str(mock_log.call_args),
@@ -1948,7 +1948,7 @@ def _input_fn():
 
     est = estimator.EstimatorV2(model_fn=_model_fn)
     est.train(dummy_input_fn, steps=1)
-    with tf.test.mock.patch.object(logging, 'warning') as mock_log:
+    with tf.test.mock.patch.object(tf.logging, 'warning') as mock_log:
       next(est.predict(_input_fn))
       # The warning should not have keyword QueueRunner.
       self.assertRegexpMatches(str(mock_log.call_args), '^((?!QueueRunner).)*$')
@@ -1971,7 +1971,7 @@ def _input_fn():
 
     est = estimator.EstimatorV2(model_fn=_model_fn)
     est.train(dummy_input_fn, steps=1)
-    with tf.test.mock.patch.object(logging, 'warning') as mock_log:
+    with tf.test.mock.patch.object(tf.logging, 'warning') as mock_log:
       next(est.predict(_input_fn))
       # The warning should not have keyword QueueRunner.
       self.assertRegexpMatches(str(mock_log.call_args), '^((?!QueueRunner).)*$')

tensorflow_estimator/python/estimator/training.py

@@ -26,7 +26,6 @@
 import six
 import tensorflow as tf
 from tensorflow.python.distribute import estimator_training as distribute_coordinator_training
-from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.training import basic_session_run_hooks
 from tensorflow.python.training import server_lib
 from tensorflow_estimator.python.estimator import estimator as estimator_lib
@@ -851,7 +850,7 @@ def _start_continuous_evaluation(self):
             tf.compat.v1.GraphKeys.GLOBAL_STEP)
         if (global_step and self._train_spec.max_steps and
             global_step >= self._train_spec.max_steps):
-          logging.info(
+          tf.compat.v1.logging.info(
               'Exiting evaluation, global_step=%s >= train max_steps=%s',
               global_step, self._train_spec.max_steps)
           return
@@ -894,11 +893,13 @@ def _execute_evaluator_once(self, evaluator, continuous_eval_listener,
     elapsed_time = time.time() - start
     difference = throttle_secs - elapsed_time
     if difference > 0:
-      logging.info('Waiting %f secs before starting next eval run.', difference)
+      tf.compat.v1.logging.info(
+          'Waiting %f secs before starting next eval run.', difference
+      )
       time.sleep(difference)
     elif (throttle_secs == 0 and eval_result.status != _EvalStatus.EVALUATED):
       # Prints a user-actionable warning to avoid unnecessary load on evaluator.
-      logging.warning(
+      tf.compat.v1.logging.warning(
           'EvalSpec.throttle_secs is set as 0. This might overload the job '
           'before finding (next) new checkpoint. Please consider to increase '
           'it.')
@@ -978,7 +979,7 @@ def _log_err_msg(self, message):
       """Prints warning `message` every 10 mins."""
       current_time = time.time()
       if current_time - self._last_warning_time > 600:
-        logging.warning(message)
+        tf.compat.v1.logging.warning(message)
         self._last_warning_time = current_time
 
     def _export_eval_result(self, eval_result, is_the_final_export):

tensorflow_estimator/python/estimator/training_test.py

@@ -28,7 +28,6 @@
 
 import numpy as np
 import tensorflow as tf
-from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.training import basic_session_run_hooks
 from tensorflow.python.training import server_lib
 from tensorflow_estimator.python.estimator import estimator as estimator_lib
@@ -1371,7 +1370,7 @@ def test_skip_evaluation_due_to_ckpt(self):
         input_fn=lambda: 1, start_delay_secs=0, throttle_secs=2)
 
     executor = training._TrainingExecutor(mock_est, mock_train_spec, eval_spec)
-    with tf.compat.v1.test.mock.patch.object(logging, 'warning') as mock_log:
+    with tf.compat.v1.test.mock.patch.object(tf.compat.v1.logging, 'warning') as mock_log:
       executor.run_evaluator()
 
     # Three checkpoint paths are invalid.
@@ -1399,7 +1398,7 @@ def test_warning_if_throttle_secs_is_zero(self):
         input_fn=lambda: 1, start_delay_secs=0, throttle_secs=0)
 
     executor = training._TrainingExecutor(mock_est, mock_train_spec, eval_spec)
-    with tf.compat.v1.test.mock.patch.object(logging, 'warning') as mock_log:
+    with tf.compat.v1.test.mock.patch.object(tf.compat.v1.logging, 'warning') as mock_log:
       executor.run_evaluator()
 
     # First ckpt is invalid.
@@ -1513,7 +1512,7 @@ def test_throttle_secs(self, mock_sleep, mock_time):
 
     executor = training._TrainingExecutor(mock_est, mock_train_spec, eval_spec)
     # Disable logging as it calls time.time also.
-    with tf.compat.v1.test.mock.patch.object(logging, 'info'):
+    with tf.compat.v1.test.mock.patch.object(tf.compat.v1.logging, 'info'):
       executor.run_evaluator()
     mock_sleep.assert_called_with(throttle_secs - operation_secs)
     self.assertTrue(mock_est.evaluate.called)