upload works (too fast?)

Author: bileschi

tensorboard/uploader/orchestration/batched_request_sender.py

@@ -189,9 +189,13 @@ def send_requests(self, run_to_events):
                     run_name, event, value, metadata
                 )
 
+        # Send requests corresponding to whatever remaining events.
         self._scalar_request_sender.flush()
         self._tensor_request_sender.flush()
         self._blob_request_sender.flush()
+        # Wait for asynchronous calls to complete.
+        self._scalar_request_sender.complete_all_pending_futures()
+
 
     def _run_values(self, run_to_events):
         """Helper generator to create a single stream of work items.

tensorboard/uploader/orchestration/scalar_batched_request_sender.py

@@ -29,6 +29,8 @@
 
 logger = tb_logging.get_logger()
 
+# How long to wait for a response on a scalar write request.
+_SCALAR_WRITE_TIMEOUT_SECS = 30
 
 def _prune_empty_tags_and_runs(request):
     for (run_idx, run) in reversed(list(enumerate(request.runs))):
@@ -149,7 +151,7 @@ def flush(self):
             return
 
         self._rpc_rate_limiter.tick()
-        self.complete_grpc_futures(grpc_futures)
+        self._groom_grpc_futures()
 
         with _request_logger(
             request, request.runs
@@ -160,8 +162,8 @@ def flush(self):
 
         self._new_request()
 
-    def complete_grpc_futures(self, grpc_futures):
-        """Handle any excptions."""
+    def _groom_grpc_futures(self):
+        """Handle any excptions, remove completed futures."""
         done_futures = []
         # Check if any exceptions raised, collect indicies of futures which can
         # be removed.
@@ -172,7 +174,7 @@ def complete_grpc_futures(self, grpc_futures):
                     # WriteScalar RPCs, but if we did, it would go here.  This
                     # call to result will raise any exception caused in the
                     # gRPC call.
-                    future.result()
+                    future.result(_SCALAR_WRITE_TIMEOUT_SECS)
                     done_futures.append(i)
                 except grpc.RpcError as e:
                     if e.code() == grpc.StatusCode.NOT_FOUND:
@@ -183,6 +185,15 @@ def complete_grpc_futures(self, grpc_futures):
         for i in done_futures:
             self._grpc_futures.pop(i)
 
+    def complete_all_pending_futures(self):
+        """Continuously checks the futures until they are done.
+
+        This is guaranteed to complete if the underlying gRPC future
+        requests are made with timeouts.
+        """
+        while self._grpc_futures:
+            self._groom_grpc_futures()
+
 
     def _create_run(self, run_name):
         """Adds a run to the live request, if there's space.

tensorboard/uploader/orchestration/scalar_batched_request_sender_test.py

@@ -57,7 +57,8 @@ def _apply_compat(events):
             yield event
 
 
-def _create_mock_client():
+
+def _create_mock_client(grpc_exception=None):
     # Create a stub instance (using a test channel) in order to derive a mock
     # from it with autospec enabled. Mocking TensorBoardWriterServiceStub itself
     # doesn't work with autospec because grpc constructs stubs via metaclassing.
@@ -66,6 +67,30 @@ def _create_mock_client():
     )
     stub = write_service_pb2_grpc.TensorBoardWriterServiceStub(test_channel)
     mock_client = mock.create_autospec(stub)
+
+    # Some surgery is required to make the mock object returned from the call to
+    # gRPC's `future` async endpoint act like the `grpc.Future` it is supposed
+    # to be. It needs to somewhat faithfully reproduce the following methods
+    # which are used in the grpc_util.async_call_with_retries machinery.
+    #  *  `done`
+    #  *  `result`
+    #  *  `exception`
+    grpc_future = mock.MagicMock()
+    grpc_future.done.return_value = True
+    if grpc_exception:
+        grpc_future.result.side_effect = grpc_exception
+        grpc_future.exception.return_value = grpc_exception
+    else:
+        grpc_future.exception.return_value = None
+    mock_client.WriteScalar.future.return_value = grpc_future
+
+    # This is reimplementing the gRPC future behavior wherein a callback
+    # can be invoked when the future is complete.  In this case, the callback
+    # must be invoked with the grpc_future as its argument.
+    def _invoke_callback(handler):
+        handler(grpc_future)
+
+    grpc_future.add_done_callback.side_effect = _invoke_callback
     return mock_client
 
 
@@ -117,8 +142,9 @@ def _add_events_and_flush(self, events):
         )
         self._add_events(sender, "", events)
         sender.flush()
+        sender.complete_all_pending_futures()
 
-        requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
+        requests = [c[0][0] for c in mock_client.WriteScalar.future.call_args_list]
         self.assertLen(requests, 1)
         self.assertLen(requests[0].runs, 1)
         return requests[0].runs[0]
@@ -205,18 +231,22 @@ def test_v2_summary(self):
         )
         self.assertProtoEquals(run_proto, expected_run_proto)
 
-    def test_propagates_experiment_deletion(self):
+    def test_async_propagates_experiment_deletion(self):
+        # Setup: A client which return a grpc_error when asked to write.
         event = event_pb2.Event(step=1)
         event.summary.value.add(tag="foo", simple_value=1.0)
 
-        mock_client = _create_mock_client()
+        error = test_util.grpc_error(grpc.StatusCode.NOT_FOUND, "nope")
+        mock_client = _create_mock_client(grpc_exception=error)
         sender = _create_scalar_request_sender("123", mock_client)
+        # Execute: sender.add_event on all the events.
         self._add_events(sender, "run", _apply_compat([event]))
-
-        error = test_util.grpc_error(grpc.StatusCode.NOT_FOUND, "nope")
-        mock_client.WriteScalar.side_effect = error
+        # Expect: Exception is raised, but it is the transforemd
+        # uploader exception, not the original grpc exception.
         with self.assertRaises(uploader_errors.ExperimentNotFoundError):
             sender.flush()
+            sender.complete_all_pending_futures()
+
 
     def test_no_budget_for_base_request(self):
         mock_client = _create_mock_client()
@@ -263,7 +293,8 @@ def test_break_at_run_boundary(self):
         self._add_events(sender, long_run_1, _apply_compat([event_1]))
         self._add_events(sender, long_run_2, _apply_compat([event_2]))
         sender.flush()
-        requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
+        sender.complete_all_pending_futures()
+        requests = [c[0][0] for c in mock_client.WriteScalar.future.call_args_list]
 
         for request in requests:
             _clear_wall_times(request)
@@ -306,7 +337,8 @@ def test_break_at_tag_boundary(self):
         )
         self._add_events(sender, "train", _apply_compat([event]))
         sender.flush()
-        requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
+        sender.complete_all_pending_futures()
+        requests = [c[0][0] for c in mock_client.WriteScalar.future.call_args_list]
         for request in requests:
             _clear_wall_times(request)
 
@@ -352,7 +384,8 @@ def test_break_at_scalar_point_boundary(self):
         )
         self._add_events(sender, "train", _apply_compat(events))
         sender.flush()
-        requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
+        sender.complete_all_pending_futures()
+        requests = [c[0][0] for c in mock_client.WriteScalar.future.call_args_list]
         for request in requests:
             _clear_wall_times(request)
 
@@ -407,7 +440,8 @@ def mock_add_point(byte_budget_manager_self, point):
             self._add_events(sender, "train", _apply_compat([event_1]))
             self._add_events(sender, "test", _apply_compat([event_2]))
             sender.flush()
-        requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list]
+            sender.complete_all_pending_futures()
+        requests = [c[0][0] for c in mock_client.WriteScalar.future.call_args_list]
         for request in requests:
             _clear_wall_times(request)