Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 56 additions & 53 deletions tensorboard/uploader/uploader_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,21 +68,55 @@ def _create_mock_client():
return mock_client


# Sentinel for `_create_*` helpers, for arguments for which we want to
# supply a default other than the `None` used by the code under test.
_USE_DEFAULT = object()


def _create_uploader(
writer_client=_USE_DEFAULT,
logdir=None,
rpc_rate_limiter=_USE_DEFAULT,
name=None,
description=None,
):
if writer_client is _USE_DEFAULT:
writer_client = _create_mock_client()
if rpc_rate_limiter is _USE_DEFAULT:
rpc_rate_limiter = util.RateLimiter(0)
return uploader_lib.TensorBoardUploader(
writer_client,
logdir,
rpc_rate_limiter=rpc_rate_limiter,
name=name,
description=description,
)


def _create_request_sender(
experiment_id=None, api=None, rpc_rate_limiter=_USE_DEFAULT
):
if api is _USE_DEFAULT:
api = _create_mock_client()
if rpc_rate_limiter is _USE_DEFAULT:
rpc_rate_limiter = util.RateLimiter(0)
return uploader_lib._BatchedRequestSender(
experiment_id=experiment_id, api=api, rpc_rate_limiter=rpc_rate_limiter
)


class TensorboardUploaderTest(tf.test.TestCase):
def test_create_experiment(self):
logdir = "/logs/foo"
mock_client = _create_mock_client()
uploader = uploader_lib.TensorBoardUploader(mock_client, logdir)
uploader = _create_uploader(_create_mock_client(), logdir)
eid = uploader.create_experiment()
self.assertEqual(eid, "123")

def test_create_experiment_with_name(self):
logdir = "/logs/foo"
mock_client = _create_mock_client()
new_name = "This is the new name"
uploader = uploader_lib.TensorBoardUploader(
mock_client, logdir, name=new_name
)
uploader = _create_uploader(mock_client, logdir, name=new_name)
eid = uploader.create_experiment()
self.assertEqual(eid, "123")
mock_client.CreateExperiment.assert_called_once()
Expand All @@ -100,7 +134,7 @@ def test_create_experiment_with_description(self):
**description**"
may have "strange" unicode chars 🌴 \\/<>
"""
uploader = uploader_lib.TensorBoardUploader(
uploader = _create_uploader(
mock_client, logdir, description=new_description
)
eid = uploader.create_experiment()
Expand All @@ -121,7 +155,7 @@ def test_create_experiment_with_all_metadata(self):
may have "strange" unicode chars 🌴 \/<>
"""
new_name = "This is a cool name."
uploader = uploader_lib.TensorBoardUploader(
uploader = _create_uploader(
mock_client, logdir, name=new_name, description=new_description
)
eid = uploader.create_experiment()
Expand All @@ -136,7 +170,7 @@ def test_create_experiment_with_all_metadata(self):

def test_start_uploading_without_create_experiment_fails(self):
mock_client = _create_mock_client()
uploader = uploader_lib.TensorBoardUploader(mock_client, "/logs/foo")
uploader = _create_uploader(mock_client, "/logs/foo")
with self.assertRaisesRegex(RuntimeError, "call create_experiment()"):
uploader.start_uploading()

Expand All @@ -145,8 +179,8 @@ def test_start_uploading_without_create_experiment_fails(self):
def test_start_uploading(self):
mock_client = _create_mock_client()
mock_rate_limiter = mock.create_autospec(util.RateLimiter)
uploader = uploader_lib.TensorBoardUploader(
mock_client, "/logs/foo", mock_rate_limiter
uploader = _create_uploader(
mock_client, "/logs/foo", rpc_rate_limiter=mock_rate_limiter
)
uploader.create_experiment()

Expand Down Expand Up @@ -177,10 +211,7 @@ def scalar_event(tag, value):
def test_upload_empty_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_uploader(mock_client, logdir)
uploader.create_experiment()
uploader._upload_once()
mock_client.WriteScalar.assert_not_called()
Expand All @@ -190,10 +221,7 @@ def test_upload_swallows_rpc_failure(self):
with tb_test_util.FileWriter(logdir) as writer:
writer.add_test_summary("foo")
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_uploader(mock_client, logdir)
uploader.create_experiment()
error = test_util.grpc_error(grpc.StatusCode.INTERNAL, "Failure")
mock_client.WriteScalar.side_effect = error
Expand All @@ -205,10 +233,7 @@ def test_upload_propagates_experiment_deletion(self):
with tb_test_util.FileWriter(logdir) as writer:
writer.add_test_summary("foo")
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_uploader(mock_client, logdir)
uploader.create_experiment()
error = test_util.grpc_error(grpc.StatusCode.NOT_FOUND, "nope")
mock_client.WriteScalar.side_effect = error
Expand All @@ -227,10 +252,7 @@ def test_upload_preserves_wall_time(self):
)
)
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_uploader(mock_client, logdir)
uploader.create_experiment()
uploader._upload_once()
mock_client.WriteScalar.assert_called_once()
Expand All @@ -245,10 +267,7 @@ def test_upload_preserves_wall_time(self):
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_uploader(mock_client, logdir)
uploader.create_experiment()

# Convenience helpers for constructing expected requests.
Expand Down Expand Up @@ -359,11 +378,7 @@ def test_upload_full_logdir(self):
class BatchedRequestSenderTest(tf.test.TestCase):
def _populate_run_from_events(self, run_proto, events):
mock_client = _create_mock_client()
builder = uploader_lib._BatchedRequestSender(
experiment_id="123",
api=mock_client,
rpc_rate_limiter=util.RateLimiter(0),
)
builder = _create_request_sender(experiment_id="123", api=mock_client)
builder.send_requests({"": events})
requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
if requests:
Expand Down Expand Up @@ -549,9 +564,7 @@ def test_no_budget_for_experiment_id(self):
long_experiment_id = "A" * uploader_lib._MAX_REQUEST_LENGTH_BYTES
mock_client = _create_mock_client()
with self.assertRaises(RuntimeError) as cm:
builder = uploader_lib._BatchedRequestSender(
long_experiment_id, mock_client, util.RateLimiter(0)
)
builder = _create_request_sender(long_experiment_id, mock_client)
builder.send_requests(run_to_events)
self.assertEqual(
str(cm.exception), "Byte budget too small for experiment ID"
Expand All @@ -564,9 +577,7 @@ def test_no_room_for_single_point(self):
long_run_name = "A" * uploader_lib._MAX_REQUEST_LENGTH_BYTES
run_to_events = {long_run_name: [event]}
with self.assertRaises(RuntimeError) as cm:
builder = uploader_lib._BatchedRequestSender(
"123", mock_client, util.RateLimiter(0)
)
builder = _create_request_sender("123", mock_client)
builder.send_requests(run_to_events)
self.assertEqual(str(cm.exception), "add_event failed despite flush")

Expand All @@ -584,9 +595,7 @@ def test_break_at_run_boundary(self):
[(long_run_1, [event_1]), (long_run_2, [event_2])]
)

builder = uploader_lib._BatchedRequestSender(
"123", mock_client, util.RateLimiter(0)
)
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]

Expand Down Expand Up @@ -623,9 +632,7 @@ def test_break_at_tag_boundary(self):
event.summary.value.add(tag=long_tag_2, simple_value=2.0)
run_to_events = {"train": [event]}

builder = uploader_lib._BatchedRequestSender(
"123", mock_client, util.RateLimiter(0)
)
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:
Expand Down Expand Up @@ -665,9 +672,7 @@ def test_break_at_scalar_point_boundary(self):
events.append(event_pb2.Event(summary=summary, step=step))
run_to_events = {"train": events}

builder = uploader_lib._BatchedRequestSender(
"123", mock_client, util.RateLimiter(0)
)
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:
Expand Down Expand Up @@ -722,9 +727,7 @@ def mock_create_point(uploader_self, *args, **kwargs):
"_create_point",
mock_create_point,
):
builder = uploader_lib._BatchedRequestSender(
"123", mock_client, util.RateLimiter(0)
)
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:
Expand Down