def test_fails_nicely_on_stream_experiment_data_timeout(self):
# Setup: Client where:
# 1. stream_experiments will say there is one experiment_id.
# 2. stream_experiment_data will raise a grpc CANCELLED, as per
# a timeout.
mock_api_client = self._create_mock_api_client()
experiment_id = "123"
def stream_experiments(request, **kwargs):
del request # unused
yield _make_experiments_response([experiment_id])
def stream_experiment_data(request, **kwargs):
raise test_util.grpc_error(grpc.StatusCode.CANCELLED,
"details string")
mock_api_client.StreamExperiments = stream_experiments
mock_api_client.StreamExperimentData = stream_experiment_data
outdir = os.path.join(self.get_temp_dir(), "outdir")
# Execute: exporter.export()
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
generator = exporter.export()
# Expect: A nice exception of the right type and carrying the right
# experiment_id.
with self.assertRaises(exporter_lib.GrpcTimeoutException) as cm:
next(generator)
self.assertEquals(cm.exception.experiment_id, experiment_id)
def test_stream_experiment_data_passes_through_unexpected_exception(self):
# Setup: Client where:
# 1. stream_experiments will say there is one experiment_id.
# 2. stream_experiment_data will throw an internal error.
mock_api_client = self._create_mock_api_client()
experiment_id = "123"
def stream_experiments(request, **kwargs):
del request # unused
yield _make_experiments_response([experiment_id])
def stream_experiment_data(request, **kwargs):
del request # unused
raise test_util.grpc_error(grpc.StatusCode.INTERNAL,
"details string")
mock_api_client.StreamExperiments = stream_experiments
mock_api_client.StreamExperimentData = stream_experiment_data
outdir = os.path.join(self.get_temp_dir(), "outdir")
# Execute: exporter.export().
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
generator = exporter.export()
# Expect: The internal error is passed through.
with self.assertRaises(grpc.RpcError) as cm:
next(generator)
self.assertEquals(cm.exception.details(), "details string")
def test_propagates_mkdir_errors(self):
mock_api_client = self._create_mock_api_client()
outdir = os.path.join(self.get_temp_dir(), "some_file", "outdir")
with open(os.path.join(self.get_temp_dir(), "some_file"), "w"):
pass
with self.assertRaises(OSError):
exporter_lib.TensorBoardExporter(mock_api_client, outdir)
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
def test_rejects_existing_directory(self):
mock_api_client = self._create_mock_api_client()
outdir = os.path.join(self.get_temp_dir(), "outdir")
os.mkdir(outdir)
with open(os.path.join(outdir, "scalars_999.json"), "w"):
pass
with self.assertRaises(exporter_lib.OutputDirectoryExistsError):
exporter_lib.TensorBoardExporter(mock_api_client, outdir)
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
def execute(self, channel):
api_client = export_service_pb2_grpc.TensorBoardExporterServiceStub(channel)
outdir = self.output_dir
try:
exporter = exporter_lib.TensorBoardExporter(api_client, outdir)
except exporter_lib.OutputDirectoryExistsError:
msg = 'Output directory already exists: %r' % outdir
raise base_plugin.FlagsError(msg)
num_experiments = 0
for experiment_id in exporter.export():
num_experiments += 1
print('Downloaded experiment %s' % experiment_id)
print('Done. Downloaded %d experiments to: %s' % (num_experiments, outdir))
def test_handles_outdir_with_no_slash(self):
oldcwd = os.getcwd()
try:
os.chdir(self.get_temp_dir())
mock_api_client = self._create_mock_api_client()
mock_api_client.StreamExperiments.return_value = iter(
[_make_experiments_response(["123"])])
mock_api_client.StreamExperimentData.return_value = iter(
[export_service_pb2.StreamExperimentDataResponse()])
exporter = exporter_lib.TensorBoardExporter(
mock_api_client, "outdir")
generator = exporter.export()
self.assertEqual(list(generator), ["123"])
self.assertTrue(os.path.isdir("outdir"))
finally:
os.chdir(oldcwd)
def test_rejects_existing_file(self):
mock_api_client = self._create_mock_api_client()
def stream_experiments(request, **kwargs):
del request # unused
yield export_service_pb2.StreamExperimentsResponse(experiment_ids=["123"])
mock_api_client.StreamExperiments = stream_experiments
outdir = os.path.join(self.get_temp_dir(), "outdir")
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
generator = exporter.export()
with open(os.path.join(outdir, "scalars_123.json"), "w"):
pass
with self.assertRaises(exporter_lib.OutputFileExistsError):
next(generator)
mock_api_client.StreamExperimentData.assert_not_called()
def test_rejects_dangerous_experiment_ids(self):
mock_api_client = self._create_mock_api_client()
def stream_experiments(request, **kwargs):
del request # unused
yield _make_experiments_response(["../authorized_keys"])
mock_api_client.StreamExperiments = stream_experiments
outdir = os.path.join(self.get_temp_dir(), "outdir")
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
generator = exporter.export()
with self.assertRaises(RuntimeError) as cm:
next(generator)
msg = str(cm.exception)
self.assertIn("Unexpected characters", msg)
self.assertIn(repr(sorted([".", "/"])), msg)
self.assertIn("../authorized_keys", msg)
mock_api_client.StreamExperimentData.assert_not_called()
def execute(self, server_info, channel):
api_client = export_service_pb2_grpc.TensorBoardExporterServiceStub(
channel)
outdir = self.output_dir
try:
exporter = exporter_lib.TensorBoardExporter(api_client, outdir)
except exporter_lib.OutputDirectoryExistsError:
msg = "Output directory already exists: %r" % outdir
raise base_plugin.FlagsError(msg)
num_experiments = 0
try:
for experiment_id in exporter.export():
num_experiments += 1
print("Downloaded experiment %s" % experiment_id)
except exporter_lib.GrpcTimeoutException as e:
print(
"\nUploader has failed because of a timeout error. Please reach "
"out via e-mail to [email protected] to get help "
"completing your export of experiment %s." % e.experiment_id)
print("Done. Downloaded %d experiments to: %s" %
(num_experiments, outdir))
def test_e2e_success_case_with_only_scalar_data(self):
mock_api_client = self._create_mock_api_client()
mock_api_client.StreamExperiments.return_value = iter(
[_make_experiments_response(["789"])])
def stream_experiments(request, **kwargs):
del request # unused
self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
response = export_service_pb2.StreamExperimentsResponse()
response.experiments.add(experiment_id="123")
response.experiments.add(experiment_id="456")
yield response
response = export_service_pb2.StreamExperimentsResponse()
experiment = response.experiments.add()
experiment.experiment_id = "789"
experiment.name = "bert"
experiment.description = "ernie"
util.set_timestamp(experiment.create_time, 981173106)
util.set_timestamp(experiment.update_time, 1015218367)
yield response
def stream_experiment_data(request, **kwargs):
self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
for run in ("train", "test"):
for tag in ("accuracy", "loss"):
response = export_service_pb2.StreamExperimentDataResponse(
)
response.run_name = run
response.tag_name = tag
display_name = "%s:%s" % (request.experiment_id, tag)
response.tag_metadata.CopyFrom(
test_util.scalar_metadata(display_name))
for step in range(10):
response.points.steps.append(step)
response.points.values.append(2.0 * step)
response.points.wall_times.add(seconds=1571084520 +
step,
nanos=862939144)
yield response
mock_api_client.StreamExperiments = mock.Mock(wraps=stream_experiments)
mock_api_client.StreamExperimentData = mock.Mock(
wraps=stream_experiment_data)
outdir = os.path.join(self.get_temp_dir(), "outdir")
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
start_time = 1571084846.25
start_time_pb = test_util.timestamp_pb(1571084846250000000)
generator = exporter.export(read_time=start_time)
expected_files = []
self.assertTrue(os.path.isdir(outdir))
self.assertCountEqual(expected_files, _outdir_files(outdir))
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
# The first iteration should request the list of experiments and
# data for one of them.
self.assertEqual(next(generator), "123")
expected_files.append(os.path.join("experiment_123", "metadata.json"))
expected_files.append(os.path.join("experiment_123", "scalars.json"))
expected_files.append(os.path.join("experiment_123", "tensors.json"))
# blob_sequences.json should exist and be empty.
expected_files.append(
os.path.join("experiment_123", "blob_sequences.json"))
self.assertCountEqual(expected_files, _outdir_files(outdir))
# Check that the tensors and blob_sequences data files are empty, because
# there are no tensors or blob sequences.
with open(os.path.join(outdir, "experiment_123",
"tensors.json")) as infile:
self.assertEqual(infile.read(), "")
with open(os.path.join(outdir, "experiment_123",
"blob_sequences.json")) as infile:
self.assertEqual(infile.read(), "")
expected_eids_request = export_service_pb2.StreamExperimentsRequest()
expected_eids_request.read_timestamp.CopyFrom(start_time_pb)
expected_eids_request.limit = 2**63 - 1
expected_eids_request.experiments_mask.create_time = True
expected_eids_request.experiments_mask.update_time = True
expected_eids_request.experiments_mask.name = True
expected_eids_request.experiments_mask.description = True
mock_api_client.StreamExperiments.assert_called_once_with(
expected_eids_request, metadata=grpc_util.version_metadata())
expected_data_request = export_service_pb2.StreamExperimentDataRequest(
)
expected_data_request.experiment_id = "123"
expected_data_request.read_timestamp.CopyFrom(start_time_pb)
mock_api_client.StreamExperimentData.assert_called_once_with(
expected_data_request, metadata=grpc_util.version_metadata())
# The next iteration should just request data for the next experiment.
mock_api_client.StreamExperiments.reset_mock()
mock_api_client.StreamExperimentData.reset_mock()
self.assertEqual(next(generator), "456")
expected_files.append(os.path.join("experiment_456", "metadata.json"))
expected_files.append(os.path.join("experiment_456", "scalars.json"))
expected_files.append(os.path.join("experiment_456", "tensors.json"))
# blob_sequences.json should exist and be empty.
expected_files.append(
os.path.join("experiment_456", "blob_sequences.json"))
self.assertCountEqual(expected_files, _outdir_files(outdir))
mock_api_client.StreamExperiments.assert_not_called()
expected_data_request.experiment_id = "456"
mock_api_client.StreamExperimentData.assert_called_once_with(
expected_data_request, metadata=grpc_util.version_metadata())
# Again, request data for the next experiment; this experiment ID
# was in the second response batch in the list of IDs.
expected_files.append(os.path.join("experiment_789", "metadata.json"))
expected_files.append(os.path.join("experiment_789", "scalars.json"))
expected_files.append(os.path.join("experiment_789", "tensors.json"))
# blob_sequences.json should exist and be empty.
expected_files.append(
os.path.join("experiment_789", "blob_sequences.json"))
mock_api_client.StreamExperiments.reset_mock()
mock_api_client.StreamExperimentData.reset_mock()
self.assertEqual(next(generator), "789")
self.assertCountEqual(expected_files, _outdir_files(outdir))
mock_api_client.StreamExperiments.assert_not_called()
expected_data_request.experiment_id = "789"
mock_api_client.StreamExperimentData.assert_called_once_with(
expected_data_request, metadata=grpc_util.version_metadata())
# The final continuation shouldn't need to send any RPCs.
mock_api_client.StreamExperiments.reset_mock()
mock_api_client.StreamExperimentData.reset_mock()
self.assertEqual(list(generator), [])
self.assertCountEqual(expected_files, _outdir_files(outdir))
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
# Spot-check one of the scalar data files.
with open(os.path.join(outdir, "experiment_456",
"scalars.json")) as infile:
jsons = [json.loads(line) for line in infile]
self.assertLen(jsons, 4)
datum = jsons[2]
self.assertEqual(datum.pop("run"), "test")
self.assertEqual(datum.pop("tag"), "accuracy")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
expected_summary_metadata = test_util.scalar_metadata("456:accuracy")
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
expected_steps = [x for x in range(10)]
expected_values = [2.0 * x for x in range(10)]
expected_wall_times = [1571084520.862939144 + x for x in range(10)]
self.assertEqual(points.pop("steps"), expected_steps)
self.assertEqual(points.pop("values"), expected_values)
self.assertEqual(points.pop("wall_times"), expected_wall_times)
self.assertEqual(points, {})
self.assertEqual(datum, {})
# Check that one of the blob_sequences data file is empty, because there
# no blob sequences in this experiment.
with open(os.path.join(outdir, "experiment_456",
"blob_sequences.json")) as infile:
self.assertEqual(infile.read(), "")
# Spot-check one of the metadata files.
with open(os.path.join(outdir, "experiment_789",
"metadata.json")) as infile:
metadata = json.load(infile)
self.assertEqual(
metadata,
{
"name": "bert",
"description": "ernie",
"create_time": "2001-02-03T04:05:06Z",
"update_time": "2002-03-04T05:06:07Z",
},
)
def test_e2e_success_case_with_blob_sequence_data(self):
"""Covers exporting of complete and incomplete blob sequences
as well as rpc error during blob streaming.
"""
mock_api_client = self._create_mock_api_client()
def stream_experiments(request, **kwargs):
del request # unused
self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
response = export_service_pb2.StreamExperimentsResponse()
response.experiments.add(experiment_id="123")
yield response
response = export_service_pb2.StreamExperimentsResponse()
response.experiments.add(experiment_id="456")
yield response
def stream_experiment_data(request, **kwargs):
self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
tag = "__default_graph__"
for run in ("train", "test"):
response = export_service_pb2.StreamExperimentDataResponse()
response.run_name = run
response.tag_name = tag
display_name = "%s:%s" % (request.experiment_id, tag)
response.tag_metadata.CopyFrom(
summary_pb2.SummaryMetadata(
data_class=summary_pb2.DATA_CLASS_BLOB_SEQUENCE))
for step in range(1):
response.blob_sequences.steps.append(step)
response.blob_sequences.wall_times.add(seconds=1571084520 +
step,
nanos=862939144)
blob_sequence = blob_pb2.BlobSequence()
if run == "train":
# A finished blob sequence.
blob = blob_pb2.Blob(
blob_id="%s_blob" % run,
state=blob_pb2.BlobState.BLOB_STATE_CURRENT,
)
blob_sequence.entries.append(
blob_pb2.BlobSequenceEntry(blob=blob))
# An unfinished blob sequence.
blob = blob_pb2.Blob(
state=blob_pb2.BlobState.BLOB_STATE_UNFINALIZED, )
blob_sequence.entries.append(
blob_pb2.BlobSequenceEntry(blob=blob))
elif run == "test":
blob_sequence.entries.append(
# `blob` unspecified: a hole in the blob sequence.
blob_pb2.BlobSequenceEntry())
response.blob_sequences.values.append(blob_sequence)
yield response
mock_api_client.StreamExperiments = mock.Mock(wraps=stream_experiments)
mock_api_client.StreamExperimentData = mock.Mock(
wraps=stream_experiment_data)
mock_api_client.StreamBlobData.side_effect = [
iter([
export_service_pb2.StreamBlobDataResponse(
data=b"4321",
offset=0,
final_chunk=False,
),
export_service_pb2.StreamBlobDataResponse(
data=b"8765",
offset=4,
final_chunk=True,
),
]),
# Raise error from `StreamBlobData` to test the grpc-error
# condition.
test_util.grpc_error(grpc.StatusCode.INTERNAL,
"Error for testing"),
]
outdir = os.path.join(self.get_temp_dir(), "outdir")
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
start_time = 1571084846.25
start_time_pb = test_util.timestamp_pb(1571084846250000000)
generator = exporter.export(read_time=start_time)
expected_files = []
self.assertTrue(os.path.isdir(outdir))
self.assertCountEqual(expected_files, _outdir_files(outdir))
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
# The first iteration should request the list of experiments and
# data for one of them.
self.assertEqual(next(generator), "123")
expected_files.append(os.path.join("experiment_123", "metadata.json"))
# scalars.json and tensors.json should exist and be empty.
expected_files.append(os.path.join("experiment_123", "scalars.json"))
expected_files.append(os.path.join("experiment_123", "tensors.json"))
expected_files.append(
os.path.join("experiment_123", "blob_sequences.json"))
expected_files.append(
os.path.join("experiment_123", "blobs", "blob_train_blob.bin"))
# blobs/blob_test_blob.bin should not exist, because it contains
# an unfinished blob.
self.assertCountEqual(expected_files, _outdir_files(outdir))
# Check that the scalars and tensors data files are empty, because there
# no scalars or tensors.
with open(os.path.join(outdir, "experiment_123",
"scalars.json")) as infile:
self.assertEqual(infile.read(), "")
with open(os.path.join(outdir, "experiment_123",
"tensors.json")) as infile:
self.assertEqual(infile.read(), "")
# Check the blob_sequences.json file.
with open(os.path.join(outdir, "experiment_123",
"blob_sequences.json")) as infile:
jsons = [json.loads(line) for line in infile]
self.assertLen(jsons, 2)
datum = jsons[0]
self.assertEqual(datum.pop("run"), "train")
self.assertEqual(datum.pop("tag"), "__default_graph__")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
expected_summary_metadata = summary_pb2.SummaryMetadata(
data_class=summary_pb2.DATA_CLASS_BLOB_SEQUENCE)
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
self.assertEqual(datum, {})
self.assertEqual(points.pop("steps"), [0])
self.assertEqual(points.pop("wall_times"), [1571084520.862939144])
# The 1st blob is finished; the 2nd is unfinished.
self.assertEqual(points.pop("blob_file_paths"),
[["blobs/blob_train_blob.bin", None]])
self.assertEqual(points, {})
datum = jsons[1]
self.assertEqual(datum.pop("run"), "test")
self.assertEqual(datum.pop("tag"), "__default_graph__")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
self.assertEqual(datum, {})
self.assertEqual(points.pop("steps"), [0])
self.assertEqual(points.pop("wall_times"), [1571084520.862939144])
# `None` blob file path indicates an unfinished blob.
self.assertEqual(points.pop("blob_file_paths"), [[None]])
self.assertEqual(points, {})
# Check the BLOB files.
with open(
os.path.join(outdir, "experiment_123", "blobs",
"blob_train_blob.bin"),
"rb",
) as f:
self.assertEqual(f.read(), b"43218765")
# Check call to StreamBlobData.
expected_blob_data_request = export_service_pb2.StreamBlobDataRequest(
blob_id="train_blob")
mock_api_client.StreamBlobData.assert_called_once_with(
expected_blob_data_request, metadata=grpc_util.version_metadata())
# Test the case where blob streaming errors out.
self.assertEqual(next(generator), "456")
# Check the blob_sequences.json file.
with open(os.path.join(outdir, "experiment_456",
"blob_sequences.json")) as infile:
jsons = [json.loads(line) for line in infile]
self.assertLen(jsons, 2)
datum = jsons[0]
self.assertEqual(datum.pop("run"), "train")
self.assertEqual(datum.pop("tag"), "__default_graph__")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
self.assertEqual(datum, {})
self.assertEqual(points.pop("steps"), [0])
self.assertEqual(points.pop("wall_times"), [1571084520.862939144])
# `None` represents the blob that experienced error during downloading
# and hence is missing.
self.assertEqual(points.pop("blob_file_paths"), [[None, None]])
self.assertEqual(points, {})
datum = jsons[1]
self.assertEqual(datum.pop("run"), "test")
self.assertEqual(datum.pop("tag"), "__default_graph__")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
self.assertEqual(datum, {})
self.assertEqual(points.pop("steps"), [0])
self.assertEqual(points.pop("wall_times"), [1571084520.862939144])
# `None` represents the blob that experienced error during downloading
# and hence is missing.
self.assertEqual(points.pop("blob_file_paths"), [[None]])
self.assertEqual(points, {})
def test_e2e_success_case_with_only_tensors_data(self):
mock_api_client = self._create_mock_api_client()
def stream_experiments(request, **kwargs):
del request # unused
self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
response = export_service_pb2.StreamExperimentsResponse()
response.experiments.add(experiment_id="123")
yield response
def stream_experiment_data(request, **kwargs):
self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
for run in ("train_1", "train_2"):
for tag in ("dense_1/kernel", "dense_1/bias", "text/test"):
response = export_service_pb2.StreamExperimentDataResponse(
)
response.run_name = run
response.tag_name = tag
display_name = "%s:%s" % (request.experiment_id, tag)
response.tag_metadata.CopyFrom(
test_util.scalar_metadata(display_name))
for step in range(2):
response.tensors.steps.append(step)
response.tensors.wall_times.add(
seconds=1571084520 + step,
nanos=862939144 if run == "train_1" else 962939144,
)
if tag != "text/test":
response.tensors.values.append(
tensor_util.make_tensor_proto(
np.ones([3, 2]) * step))
else:
response.tensors.values.append(
tensor_util.make_tensor_proto(
np.full([3], "a" * (step + 1))))
yield response
mock_api_client.StreamExperiments = mock.Mock(wraps=stream_experiments)
mock_api_client.StreamExperimentData = mock.Mock(
wraps=stream_experiment_data)
outdir = os.path.join(self.get_temp_dir(), "outdir")
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
start_time = 1571084846.25
start_time_pb = test_util.timestamp_pb(1571084846250000000)
generator = exporter.export(read_time=start_time)
expected_files = []
self.assertTrue(os.path.isdir(outdir))
self.assertCountEqual(expected_files, _outdir_files(outdir))
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
# The first iteration should request the list of experiments and
# data for one of them.
self.assertEqual(next(generator), "123")
expected_files.append(os.path.join("experiment_123", "metadata.json"))
# scalars.json should exist and be empty.
expected_files.append(os.path.join("experiment_123", "scalars.json"))
expected_files.append(os.path.join("experiment_123", "tensors.json"))
# blob_sequences.json should exist and be empty.
expected_files.append(
os.path.join("experiment_123", "blob_sequences.json"))
expected_files.append(
os.path.join("experiment_123", "tensors", "1571084520.862939.npz"))
expected_files.append(
os.path.join("experiment_123", "tensors",
"1571084520.862939_1.npz"))
expected_files.append(
os.path.join("experiment_123", "tensors",
"1571084520.862939_2.npz"))
expected_files.append(
os.path.join("experiment_123", "tensors", "1571084520.962939.npz"))
expected_files.append(
os.path.join("experiment_123", "tensors",
"1571084520.962939_1.npz"))
expected_files.append(
os.path.join("experiment_123", "tensors",
"1571084520.962939_2.npz"))
self.assertCountEqual(expected_files, _outdir_files(outdir))
# Check that the scalars and blob_sequences data files are empty, because
# there are no scalars or blob sequences.
with open(os.path.join(outdir, "experiment_123",
"scalars.json")) as infile:
self.assertEqual(infile.read(), "")
with open(os.path.join(outdir, "experiment_123",
"blob_sequences.json")) as infile:
self.assertEqual(infile.read(), "")
expected_eids_request = export_service_pb2.StreamExperimentsRequest()
expected_eids_request.read_timestamp.CopyFrom(start_time_pb)
expected_eids_request.limit = 2**63 - 1
expected_eids_request.experiments_mask.create_time = True
expected_eids_request.experiments_mask.update_time = True
expected_eids_request.experiments_mask.name = True
expected_eids_request.experiments_mask.description = True
mock_api_client.StreamExperiments.assert_called_once_with(
expected_eids_request, metadata=grpc_util.version_metadata())
expected_data_request = export_service_pb2.StreamExperimentDataRequest(
)
expected_data_request.experiment_id = "123"
expected_data_request.read_timestamp.CopyFrom(start_time_pb)
mock_api_client.StreamExperimentData.assert_called_once_with(
expected_data_request, metadata=grpc_util.version_metadata())
# The final StreamExperiments continuation shouldn't need to send any
# RPCs.
mock_api_client.StreamExperiments.reset_mock()
mock_api_client.StreamExperimentData.reset_mock()
self.assertEqual(list(generator), [])
# Check tensor data.
with open(os.path.join(outdir, "experiment_123",
"tensors.json")) as infile:
jsons = [json.loads(line) for line in infile]
self.assertLen(jsons, 6)
datum = jsons[0]
self.assertEqual(datum.pop("run"), "train_1")
self.assertEqual(datum.pop("tag"), "dense_1/kernel")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
expected_summary_metadata = test_util.scalar_metadata(
"123:dense_1/kernel")
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
self.assertEqual(points.pop("steps"), [0, 1])
self.assertEqual(
points.pop("tensors_file_path"),
os.path.join("tensors", "1571084520.862939.npz"),
)
self.assertEqual(datum, {})
datum = jsons[4]
self.assertEqual(datum.pop("run"), "train_2")
self.assertEqual(datum.pop("tag"), "dense_1/bias")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
expected_summary_metadata = test_util.scalar_metadata(
"123:dense_1/bias")
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
self.assertEqual(points.pop("steps"), [0, 1])
self.assertEqual(
points.pop("tensors_file_path"),
os.path.join("tensors", "1571084520.962939_1.npz"),
)
self.assertEqual(datum, {})
# Load and check the tensor data from the save .npz files.
for filename in (
"1571084520.862939.npz",
"1571084520.862939_1.npz",
"1571084520.962939.npz",
"1571084520.962939_1.npz",
):
tensors = np.load(
os.path.join(outdir, "experiment_123", "tensors", filename))
tensors = [tensors[key] for key in tensors.keys()]
self.assertLen(tensors, 2)
np.testing.assert_array_equal(tensors[0], 0 * np.ones([3, 2]))
np.testing.assert_array_equal(tensors[1], 1 * np.ones([3, 2]))
for filename in (
"1571084520.862939_2.npz",
"1571084520.962939_2.npz",
):
tensors = np.load(
os.path.join(outdir, "experiment_123", "tensors", filename))
tensors = [tensors[key] for key in tensors.keys()]
self.assertLen(tensors, 2)
np.testing.assert_array_equal(tensors[0],
np.array(["a", "a", "a"], "|S"))
np.testing.assert_array_equal(tensors[1],
np.array(["aa", "aa", "aa"], "|S"))
def test_e2e_success_case(self):
mock_api_client = self._create_mock_api_client()
mock_api_client.StreamExperiments.return_value = iter([
export_service_pb2.StreamExperimentsResponse(
experiment_ids=["789"]),
])
def stream_experiments(request):
del request # unused
yield export_service_pb2.StreamExperimentsResponse(
experiment_ids=["123", "456"])
yield export_service_pb2.StreamExperimentsResponse(
experiment_ids=["789"])
def stream_experiment_data(request):
for run in ("train", "test"):
for tag in ("accuracy", "loss"):
response = export_service_pb2.StreamExperimentDataResponse(
)
response.run_name = run
response.tag_name = tag
display_name = "%s:%s" % (request.experiment_id, tag)
response.tag_metadata.CopyFrom(
test_util.scalar_metadata(display_name))
for step in range(10):
response.points.steps.append(step)
response.points.values.append(2.0 * step)
response.points.wall_times.add(seconds=1571084520 +
step,
nanos=862939144)
yield response
mock_api_client.StreamExperiments = mock.Mock(wraps=stream_experiments)
mock_api_client.StreamExperimentData = mock.Mock(
wraps=stream_experiment_data)
outdir = os.path.join(self.get_temp_dir(), "outdir")
exporter = exporter_lib.TensorBoardExporter(mock_api_client, outdir)
start_time = 1571084846.25
start_time_pb = test_util.timestamp_pb(1571084846250000000)
generator = exporter.export(read_time=start_time)
expected_files = []
self.assertTrue(os.path.isdir(outdir))
self.assertCountEqual(expected_files, os.listdir(outdir))
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
# The first iteration should request the list of experiments and
# data for one of them.
self.assertEqual(next(generator), "123")
expected_files.append("scalars_123.json")
self.assertCountEqual(expected_files, os.listdir(outdir))
expected_eids_request = export_service_pb2.StreamExperimentsRequest()
expected_eids_request.read_timestamp.CopyFrom(start_time_pb)
expected_eids_request.limit = 2**63 - 1
mock_api_client.StreamExperiments.assert_called_once_with(
expected_eids_request)
expected_data_request = export_service_pb2.StreamExperimentDataRequest(
)
expected_data_request.experiment_id = "123"
expected_data_request.read_timestamp.CopyFrom(start_time_pb)
mock_api_client.StreamExperimentData.assert_called_once_with(
expected_data_request)
# The next iteration should just request data for the next experiment.
mock_api_client.StreamExperiments.reset_mock()
mock_api_client.StreamExperimentData.reset_mock()
self.assertEqual(next(generator), "456")
expected_files.append("scalars_456.json")
self.assertCountEqual(expected_files, os.listdir(outdir))
mock_api_client.StreamExperiments.assert_not_called()
expected_data_request.experiment_id = "456"
mock_api_client.StreamExperimentData.assert_called_once_with(
expected_data_request)
# Again, request data for the next experiment; this experiment ID
# was in the second response batch in the list of IDs.
expected_files.append("scalars_789.json")
mock_api_client.StreamExperiments.reset_mock()
mock_api_client.StreamExperimentData.reset_mock()
self.assertEqual(next(generator), "789")
self.assertCountEqual(expected_files, os.listdir(outdir))
mock_api_client.StreamExperiments.assert_not_called()
expected_data_request.experiment_id = "789"
mock_api_client.StreamExperimentData.assert_called_once_with(
expected_data_request)
# The final continuation shouldn't need to send any RPCs.
mock_api_client.StreamExperiments.reset_mock()
mock_api_client.StreamExperimentData.reset_mock()
self.assertEqual(list(generator), [])
self.assertCountEqual(expected_files, os.listdir(outdir))
mock_api_client.StreamExperiments.assert_not_called()
mock_api_client.StreamExperimentData.assert_not_called()
# Spot-check one of the files.
with open(os.path.join(outdir, "scalars_456.json")) as infile:
jsons = [json.loads(line) for line in infile]
self.assertLen(jsons, 4)
datum = jsons[2]
self.assertEqual(datum.pop("run"), "test")
self.assertEqual(datum.pop("tag"), "accuracy")
summary_metadata = summary_pb2.SummaryMetadata.FromString(
base64.b64decode(datum.pop("summary_metadata")))
expected_summary_metadata = test_util.scalar_metadata("456:accuracy")
self.assertEqual(summary_metadata, expected_summary_metadata)
points = datum.pop("points")
expected_steps = [x for x in range(10)]
expected_values = [2.0 * x for x in range(10)]
expected_wall_times = [1571084520.862939144 + x for x in range(10)]
self.assertEqual(points.pop("steps"), expected_steps)
self.assertEqual(points.pop("values"), expected_values)
self.assertEqual(points.pop("wall_times"), expected_wall_times)
self.assertEqual(points, {})
self.assertEqual(datum, {})