def test_break_at_tag_boundary(self):
mock_client = _create_mock_client()
# Choose tag name sizes such that one tag fits, but not two. Note
# that tag names appear in both `Tag.name` and the summary metadata.
long_tag_1 = "a" * 384
long_tag_2 = "b" * 384
event = event_pb2.Event(step=1)
event.summary.value.add(tag=long_tag_1, simple_value=1.0)
event.summary.value.add(tag=long_tag_2, simple_value=2.0)
run_to_events = {"train": [event]}
builder = _create_request_sender("123", mock_client)
builder.send_requests(run_to_events)
requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
for request in requests:
_clear_wall_times(request)
expected = [
write_service_pb2.WriteScalarRequest(experiment_id="123"),
write_service_pb2.WriteScalarRequest(experiment_id="123"),
]
(expected[0].runs.add(name="train").tags.add(
name=long_tag_1,
metadata=test_util.scalar_metadata(long_tag_1)).points.add(
step=1, value=1.0))
(expected[1].runs.add(name="train").tags.add(
name=long_tag_2,
metadata=test_util.scalar_metadata(long_tag_2)).points.add(
step=1, value=2.0))
self.assertEqual(requests, expected)
def test_break_at_run_boundary(self):
mock_client = _create_mock_client()
# Choose run name sizes such that one run fits, but not two.
long_run_1 = "A" * 768
long_run_2 = "B" * 768
event_1 = event_pb2.Event(step=1)
event_1.summary.value.add(tag="foo", simple_value=1.0)
event_2 = event_pb2.Event(step=2)
event_2.summary.value.add(tag="bar", simple_value=-2.0)
run_to_events = collections.OrderedDict([(long_run_1, [event_1]),
(long_run_2, [event_2])])
builder = _create_request_sender("123", mock_client)
builder.send_requests(run_to_events)
requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
for request in requests:
_clear_wall_times(request)
expected = [
write_service_pb2.WriteScalarRequest(experiment_id="123"),
write_service_pb2.WriteScalarRequest(experiment_id="123"),
]
(expected[0].runs.add(name=long_run_1).tags.add(
name="foo",
metadata=test_util.scalar_metadata("foo")).points.add(step=1,
value=1.0))
(expected[1].runs.add(name=long_run_2).tags.add(
name="bar",
metadata=test_util.scalar_metadata("bar")).points.add(step=2,
value=-2.0))
self.assertEqual(requests, expected)
def test_prunes_tags_and_runs(self):
mock_client = _create_mock_client()
event_1 = event_pb2.Event(step=1)
event_1.summary.value.add(tag="foo", simple_value=1.0)
event_2 = event_pb2.Event(step=2)
event_2.summary.value.add(tag="bar", simple_value=-2.0)
run_to_events = collections.OrderedDict(
[("train", [event_1]), ("test", [event_2])]
)
real_create_point = (
uploader_lib._ScalarBatchedRequestSender._create_point
)
create_point_call_count_box = [0]
def mock_create_point(uploader_self, *args, **kwargs):
# Simulate out-of-space error the first time that we try to store
# the second point.
create_point_call_count_box[0] += 1
if create_point_call_count_box[0] == 2:
raise uploader_lib._OutOfSpaceError()
return real_create_point(uploader_self, *args, **kwargs)
with mock.patch.object(
uploader_lib._ScalarBatchedRequestSender,
"_create_point",
mock_create_point,
):
builder = uploader_lib._BatchedRequestSender(
"123", mock_client, util.RateLimiter(0)
)
builder.send_requests(run_to_events)
requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
for request in requests:
_clear_wall_times(request)
expected = [
write_service_pb2.WriteScalarRequest(experiment_id="123"),
write_service_pb2.WriteScalarRequest(experiment_id="123"),
]
(
expected[0]
.runs.add(name="train")
.tags.add(name="foo", metadata=test_util.scalar_metadata("foo"))
.points.add(step=1, value=1.0)
)
(
expected[1]
.runs.add(name="test")
.tags.add(name="bar", metadata=test_util.scalar_metadata("bar"))
.points.add(step=2, value=-2.0)
)
self.assertEqual(expected, requests)
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
def stream_experiment_data(request, **kwargs):
self.assertEqual(request.experiment_id, "789")
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)
if tag == "loss":
if run == "train":
value = 1.0 / (step + 1)
seconds = step
else:
value = -1.0 / (step + 1)
seconds = 600 + step
else: # "accuracy"
if run == "train":
value = 1.0 / (10 - step)
seconds = step * 2
else:
value = -1.0 / (10 - step)
seconds = 600 + step * 2
response.points.values.append(value)
response.points.wall_times.add(seconds=seconds,
nanos=0)
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
def test_upload_full_logdir(self):
logdir = self.get_temp_dir()
mock_client = _create_mock_client()
mock_rate_limiter = mock.create_autospec(util.RateLimiter)
uploader = uploader_lib.TensorBoardUploader(
mock_client, logdir, mock_rate_limiter
)
uploader.create_experiment()
# Convenience helpers for constructing expected requests.
run = write_service_pb2.WriteScalarRequest.Run
tag = write_service_pb2.WriteScalarRequest.Tag
point = scalar_pb2.ScalarPoint
# First round
writer = tb_test_util.FileWriter(logdir)
writer.add_test_summary("foo", simple_value=5.0, step=1)
writer.add_test_summary("foo", simple_value=6.0, step=2)
writer.add_test_summary("foo", simple_value=7.0, step=3)
writer.add_test_summary("bar", simple_value=8.0, step=3)
writer.flush()
writer_a = tb_test_util.FileWriter(os.path.join(logdir, "a"))
writer_a.add_test_summary("qux", simple_value=9.0, step=2)
writer_a.flush()
uploader._upload_once()
self.assertEqual(1, mock_client.WriteScalar.call_count)
request1 = mock_client.WriteScalar.call_args[0][0]
_clear_wall_times(request1)
expected_request1 = write_service_pb2.WriteScalarRequest(
experiment_id="123",
runs=[
run(
name=".",
tags=[
tag(
name="foo",
metadata=test_util.scalar_metadata("foo"),
points=[
point(step=1, value=5.0),
point(step=2, value=6.0),
point(step=3, value=7.0),
],
),
tag(
name="bar",
metadata=test_util.scalar_metadata("bar"),
points=[point(step=3, value=8.0)],
),
],
),
run(
name="a",
tags=[
tag(
name="qux",
metadata=test_util.scalar_metadata("qux"),
points=[point(step=2, value=9.0)],
)
],
),
],
)
self.assertProtoEquals(expected_request1, request1)
mock_client.WriteScalar.reset_mock()
# Second round
writer.add_test_summary("foo", simple_value=10.0, step=5)
writer.add_test_summary("baz", simple_value=11.0, step=1)
writer.flush()
writer_b = tb_test_util.FileWriter(os.path.join(logdir, "b"))
writer_b.add_test_summary("xyz", simple_value=12.0, step=1)
writer_b.flush()
uploader._upload_once()
self.assertEqual(1, mock_client.WriteScalar.call_count)
request2 = mock_client.WriteScalar.call_args[0][0]
_clear_wall_times(request2)
expected_request2 = write_service_pb2.WriteScalarRequest(
experiment_id="123",
runs=[
run(
name=".",
tags=[
tag(
name="foo",
metadata=test_util.scalar_metadata("foo"),
points=[point(step=5, value=10.0)],
),
tag(
name="baz",
metadata=test_util.scalar_metadata("baz"),
points=[point(step=1, value=11.0)],
),
],
),
run(
name="b",
tags=[
tag(
name="xyz",
metadata=test_util.scalar_metadata("xyz"),
points=[point(step=1, value=12.0)],
)
],
),
],
)
self.assertProtoEquals(expected_request2, request2)
mock_client.WriteScalar.reset_mock()
# Empty third round
uploader._upload_once()
mock_client.WriteScalar.assert_not_called()
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_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, {})
