Save model from run_experiment and add TODO comments.

PiperOrigin-RevId: 387643746

Author: juanuribe28

src/python/tensorflow_cloud/core/experimental/models.py

@@ -271,7 +271,8 @@ def run_experiment_cloud(run_experiment_kwargs: Dict[str, Any],
                                                           worker_config)
         run_experiment_kwargs.update(
             dict(distribution_strategy=distribution_strategy))
-        train_lib.run_experiment(**run_experiment_kwargs)
+        model, _ = train_lib.run_experiment(**run_experiment_kwargs)
+        model.save(run_experiment_kwargs['model_dir'])
 
     run_kwargs.update(dict(entry_point=None,
                            distribution_strategy=None))
@@ -282,11 +283,14 @@ def get_distribution_strategy(chief_config, worker_count, worker_config):
     """Gets a tf distribution strategy based on the cloud run config."""
     if worker_count > 0:
         if machine_config.is_tpu_config(worker_config):
-            resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+            # TODO(b/194857231) Dependency conflict for using TPUs
+            resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
+                tpu='local')
             tf.config.experimental_connect_to_cluster(resolver)
             tf.tpu.experimental.initialize_tpu_system(resolver)
             return tf.distribute.TPUStrategy(resolver)
         else:
+            # TODO(b/148619319) Saving model currently failing
             return tf.distribute.MultiWorkerMirroredStrategy()
     elif chief_config.accelerator_count > 1:
         return tf.distribute.MirroredStrategy()

src/python/tensorflow_cloud/core/experimental/tests/unit/models_test.py

@@ -77,13 +77,28 @@ def setup_run_experiment(self):
                                           mode='train_and_eval',
                                           params=config,
                                           model_dir='model_path')
-
+        self.model = mock.MagicMock()
         self.run_experiment = mock.patch.object(
             train_lib,
             'run_experiment',
             autospec=True,
+            return_value=(self.model, {})
         ).start()
 
+    def setup_tpu(self):
+        mock.patch.object(tf.tpu.experimental,
+                          'initialize_tpu_system',
+                          autospec=True).start()
+        mock.patch.object(tf.config,
+                          'experimental_connect_to_cluster',
+                          autospec=True).start()
+        mock.patch('tensorflow.distribute.cluster_resolver.TPUClusterResolver'
+                   ).start()
+        mock_tpu_strategy = mock.MagicMock(
+            spec=tf.distribute.TPUStrategy)
+        mock.patch('tensorflow.distribute.TPUStrategy',
+                   return_value=mock_tpu_strategy).start()
+
     def tearDown(self):
         mock.patch.stopall()
         super(ModelsTest, self).tearDown()
@@ -182,20 +197,8 @@ def test_run_experiment_cloud_remote(self):
         self.remote.assert_called()
         self.run_experiment.assert_called()
         self.run.assert_called()
-
-    def setup_tpu(self):
-        mock.patch.object(tf.tpu.experimental,
-                          'initialize_tpu_system',
-                          autospec=True).start()
-        mock.patch.object(tf.config,
-                          'experimental_connect_to_cluster',
-                          autospec=True).start()
-        mock.patch('tensorflow.distribute.cluster_resolver.TPUClusterResolver'
-                   ).start()
-        mock_tpu_strategy = mock.MagicMock()
-        mock_tpu_strategy.__class__ = tf.distribute.TPUStrategy
-        mock.patch('tensorflow.distribute.TPUStrategy',
-                   return_value=mock_tpu_strategy).start()
+        self.model.save.assert_called_with(
+            self.run_experiment_kwargs['model_dir'])
 
     def test_get_distribution_strategy_tpu(self):
         tpu_srategy = tf.distribute.TPUStrategy