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)
