def test_garbage_collect_exports(self): export_dir_base = tempfile.mkdtemp() + "export/" tf.compat.v1.gfile.MkDir(export_dir_base) export_dir_1 = _create_test_export_dir(export_dir_base) export_dir_2 = _create_test_export_dir(export_dir_base) export_dir_3 = _create_test_export_dir(export_dir_base) export_dir_4 = _create_test_export_dir(export_dir_base) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_1)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_2)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_3)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_4)) def _serving_input_receiver_fn(): return tf.constant([1]), None exporter = exporter_lib.LatestExporter( name="latest_exporter", serving_input_receiver_fn=_serving_input_receiver_fn, exports_to_keep=2) estimator = tf.compat.v1.test.mock.Mock(spec=estimator_lib.Estimator) # Garbage collect all but the most recent 2 exports, # where recency is determined based on the timestamp directory names. exporter.export(estimator, export_dir_base, None, None, False) self.assertFalse(tf.compat.v1.gfile.Exists(export_dir_1)) self.assertFalse(tf.compat.v1.gfile.Exists(export_dir_2)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_3)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_4))
def test_latest_exporter(self): def _serving_input_receiver_fn(): pass export_dir_base = tempfile.mkdtemp() + "export/" tf.compat.v1.gfile.MkDir(export_dir_base) exporter = exporter_lib.LatestExporter( name="latest_exporter", serving_input_receiver_fn=_serving_input_receiver_fn, assets_extra={"from/path": "to/path"}, as_text=False, exports_to_keep=5) estimator = tf.compat.v1.test.mock.Mock(spec=estimator_lib.Estimator) estimator.export_saved_model.return_value = "export_result_path" export_result = exporter.export(estimator, export_dir_base, "checkpoint_path", {}, False) self.assertEqual("export_result_path", export_result) estimator.export_saved_model.assert_called_with( export_dir_base, _serving_input_receiver_fn, assets_extra={"from/path": "to/path"}, as_text=False, checkpoint_path="checkpoint_path")
def test_error_out_if_exports_to_keep_is_zero(self): def _serving_input_receiver_fn(): pass with self.assertRaisesRegexp(ValueError, "positive number"): exporter = exporter_lib.LatestExporter( name="latest_exporter", serving_input_receiver_fn=_serving_input_receiver_fn, exports_to_keep=0) self.assertEqual("latest_exporter", exporter.name)
def test_garbage_collect_exports_with_trailing_delimiter(self): export_dir_base = tempfile.mkdtemp() + "export/" tf.compat.v1.gfile.MkDir(export_dir_base) export_dir_1 = _create_test_export_dir(export_dir_base) export_dir_2 = _create_test_export_dir(export_dir_base) export_dir_3 = _create_test_export_dir(export_dir_base) export_dir_4 = _create_test_export_dir(export_dir_base) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_1)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_2)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_3)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_4)) def _serving_input_receiver_fn(): return tf.constant([1]), None exporter = exporter_lib.LatestExporter( name="latest_exporter", serving_input_receiver_fn=_serving_input_receiver_fn, exports_to_keep=1) estimator = tf.compat.v1.test.mock.Mock(spec=estimator_lib.Estimator) # Garbage collect all but the most recent 2 exports, # where recency is determined based on the timestamp directory names. with tf.compat.v1.test.mock.patch.object( gfile, "ListDirectory") as mock_list_directory: mock_list_directory.return_value = [ os.path.basename(export_dir_1) + b"/", os.path.basename(export_dir_2) + b"/", os.path.basename(export_dir_3) + b"/", os.path.basename(export_dir_4) + b"/", ] exporter.export(estimator, export_dir_base, None, None, False) self.assertFalse(tf.compat.v1.gfile.Exists(export_dir_1)) self.assertFalse(tf.compat.v1.gfile.Exists(export_dir_2)) self.assertFalse(tf.compat.v1.gfile.Exists(export_dir_3)) self.assertTrue(tf.compat.v1.gfile.Exists(export_dir_4))
def _get_exporter(self, name, fc): feature_spec = tf.compat.v1.feature_column.make_parse_example_spec(fc) serving_input_receiver_fn = ( export_lib.build_parsing_serving_input_receiver_fn(feature_spec)) return exporter_lib.LatestExporter( name, serving_input_receiver_fn=serving_input_receiver_fn)