Remove models_entry_point dependency on tensorflow_cloud.

PiperOrigin-RevId: 389661558

Author: juanuribe28

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

@@ -16,11 +16,9 @@
 
 import os
 import pickle
-import shutil
 from typing import Any, Dict, Optional
 import uuid
 
-from . import constants
 from .. import machine_config
 from .. import run
 import tensorflow as tf
@@ -37,6 +35,10 @@
     import importlib_resources as pkg_resources
 # pylint: enable=g-import-not-at-top
 
+_PARAMS_FILE_NAME_FORMAT = '{}_params'
+_ENTRY_POINT_FORMAT = '{}.py'
+_ENTRY_POINT_TEMPLATE = 'models_entry_point.py'
+
 
 def run_models(dataset_name: str,
                model_name: str,
@@ -267,18 +269,36 @@ def run_experiment_cloud(run_experiment_kwargs: Dict[str, Any],
         dict(distribution_strategy=distribution_strategy))
     file_id = str(uuid.uuid4())
     params_file = save_params(run_experiment_kwargs, file_id)
+    entry_point = copy_entry_point(file_id, params_file)
 
-    with pkg_resources.path(__package__, 'models_entry_point.py') as path:
-        entry_point = f'{file_id}.py'
-        shutil.copyfile(str(path), entry_point)
-        run_kwargs.update(dict(entry_point=entry_point,
-                               distribution_strategy=None))
-        info = run.run(**run_kwargs)
+    run_kwargs.update(dict(entry_point=entry_point,
+                           distribution_strategy=None))
+    info = run.run(**run_kwargs)
     os.remove(entry_point)
     os.remove(params_file)
     return info
 
 
+def copy_entry_point(file_id, params_file):
+    """Copy models_entry_point and add params file name."""
+    lines = get_original_lines()
+    entry_point = _ENTRY_POINT_FORMAT.format(file_id)
+    with open(entry_point, 'w') as entry_file:
+        for line in lines:
+            if line.startswith('PARAMS_FILE_NAME = '):
+                entry_file.write(f"PARAMS_FILE_NAME = '{params_file}'\n")
+            else:
+                entry_file.write(line)
+    return entry_point
+
+
+def get_original_lines():
+    """Gets the file lines of models_entry_point.py as a list of strings."""
+    with pkg_resources.open_text(__package__, _ENTRY_POINT_TEMPLATE) as file:
+        lines = file.readlines()
+    return lines
+
+
 def get_distribution_strategy_str(run_kwargs):
     """Gets the name of a distribution strategy based on cloud run config."""
     if ('worker_count' in run_kwargs
@@ -297,7 +317,7 @@ def get_distribution_strategy_str(run_kwargs):
 
 def save_params(params, file_id):
     """Pickles the params object using the file_id as prefix."""
-    file_name = constants.PARAMS_FILE_NAME_FORMAT.format(file_id)
+    file_name = _PARAMS_FILE_NAME_FORMAT.format(file_id)
     with open(file_name, 'xb') as f:
         pickle.dump(params, f)
     return file_name

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

@@ -12,16 +12,20 @@
 # 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.
-"""Entry point file for run_experiment_cloud."""
+"""Entry point template file for run_experiment_cloud."""
 
-import os
 import pickle
 
 import tensorflow as tf
 
-from tensorflow_cloud.core.experimental import constants
 from official.core import train_lib
 
+# PARAMS_FILE_NAME provides the name of the file that cointains the
+# run_experiment_kwargs used to call run_experiment. In models.py, when copying
+# this file, the value of PARAMS_FILE_NAME is updated to contain the actual name
+# of the file.
+PARAMS_FILE_NAME = 'file_name'
+
 
 def load_params(file_name):
     with open(file_name, 'rb') as f:
@@ -50,9 +54,7 @@ def get_one_device():
 
 
 def main():
-    prefix, _ = os.path.splitext(os.path.basename(__file__))
-    file_name = constants.PARAMS_FILE_NAME_FORMAT.format(prefix)
-    run_experiment_kwargs = load_params(file_name)
+    run_experiment_kwargs = load_params(PARAMS_FILE_NAME)
     strategy_str = run_experiment_kwargs['distribution_strategy']
     strategy = _DISTRIBUTION_STRATEGIES[strategy_str]()
     run_experiment_kwargs.update(dict(

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

@@ -18,7 +18,6 @@
 import mock
 import tensorflow as tf
 
-from tensorflow_cloud.core.experimental import constants
 from tensorflow_cloud.core.experimental import models_entry_point
 from official.core import base_task
 from official.core import config_definitions
@@ -64,9 +63,8 @@ def tearDown(self):
 
     def test_main(self):
         models_entry_point.main()
-        file_name = constants.PARAMS_FILE_NAME_FORMAT.format(
-            'models_entry_point')
-        self.load_params.assert_called_with(file_name)
+
+        self.load_params.assert_called_with(models_entry_point.PARAMS_FILE_NAME)
         self.run_experiment_kwargs.update(dict(
             distribution_strategy='one_device_strategy'))
         self.run_experiment.assert_called_with(**self.run_experiment_kwargs)

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

@@ -15,27 +15,16 @@
 """Tests for the models experimental module."""
 
 import os
-import pathlib
-import shutil
 import uuid
 from absl.testing import absltest
 import mock
 
 import tensorflow as tf
 from tensorflow_cloud.core import machine_config
 from tensorflow_cloud.core import run
-from tensorflow_cloud.core.experimental import constants
 from tensorflow_cloud.core.experimental import models
 from official.vision.image_classification.efficientnet import efficientnet_model
 
-# pylint: disable=g-import-not-at-top
-try:
-    import importlib.resources as pkg_resources
-except ImportError:
-    # Backported for python<3.7
-    import importlib_resources as pkg_resources
-# pylint: enable=g-import-not-at-top
-
 
 class ModelsTest(absltest.TestCase):
 
@@ -82,11 +71,9 @@ def setup_run_models(self):
             autospec=True,
         ).start()
 
-    def setup_run_experiment_cloud(self, path):
-        self.file_id = 'test'
-
-        self.params_file = constants.PARAMS_FILE_NAME_FORMAT.format(
-            self.file_id)
+    def setup_run_experiment_cloud(self, file_id):
+        self.params_file = models._PARAMS_FILE_NAME_FORMAT.format(
+            file_id)
 
         self.save_params = mock.patch.object(
             models,
@@ -95,11 +82,11 @@ def setup_run_experiment_cloud(self, path):
             return_value=self.params_file,
         ).start()
 
-        self.path = mock.patch.object(
-            pkg_resources,
-            'path',
+        self.copy_entry_point = mock.patch.object(
+            models,
+            'copy_entry_point',
             autospec=True,
-            return_value=pathlib.Path(path),
+            return_value=models._ENTRY_POINT_FORMAT.format(file_id),
         ).start()
 
         self.remove = mock.patch.object(
@@ -111,13 +98,7 @@ def setup_run_experiment_cloud(self, path):
         self.uuid4 = mock.patch.object(
             uuid,
             'uuid4',
-            return_value=self.file_id,
-        ).start()
-
-        self.copyfile = mock.patch.object(
-            shutil,
-            'copyfile',
-            autospec=True,
+            return_value=file_id,
         ).start()
 
     def tearDown(self):
@@ -201,18 +182,45 @@ def test_run_models_remote(self):
 
     def test_run_experiment_cloud(self):
         self.setup_run(remote=False)
-        path_str = '/test'
-        self.setup_run_experiment_cloud(path_str)
+        file_id = 'test_id'
+        self.setup_run_experiment_cloud(file_id)
         run_experiment_kwargs = dict()
         models.run_experiment_cloud(
             run_experiment_kwargs=run_experiment_kwargs)
-        entry_point = f'{self.file_id}.py'
-        self.copyfile.assert_called_with(path_str, entry_point)
+        entry_point = models._ENTRY_POINT_FORMAT.format(file_id)
         self.run.assert_called_with(entry_point=entry_point,
                                     distribution_strategy=None)
         self.remove.assert_any_call(entry_point)
         self.remove.assert_any_call(self.params_file)
 
+    def setup_copy_entry_point(self):
+        self.get_original_lines = mock.patch.object(
+            models,
+            'get_original_lines',
+            autospec=True,
+            return_value=['PARAMS_FILE_NAME = not this', 'do not change'],
+        ).start()
+
+        self.open = mock.mock_open()
+        mock.patch(
+            'builtins.open',
+            self.open,
+        ).start()
+
+    def test_copy_entry_point(self):
+        self.setup_copy_entry_point()
+        file_id = 'file_id'
+        params_file = 'params_file'
+        models.copy_entry_point(file_id, params_file)
+
+        self.open.assert_called_once_with(
+            models._ENTRY_POINT_FORMAT.format(file_id),
+            'w')
+        entry_file = self.open()
+        entry_file.write.assert_any_call(
+            f"PARAMS_FILE_NAME = '{params_file}'\n")
+        entry_file.write.assert_any_call('do not change')
+
     def test_get_distribution_strategy_tpu(self):
         run_kwargs = dict(
             worker_count=1,