uploader: require experiment metadata from server (#3310)

Summary:
The `StreamExperiments` RPC response used to send only `experiment_ids`,
but now sends `experiments` with additional metadata. This server-side
change has been live since late November 2019, so we’re confident that
we won’t need to roll it back. Thus, we can drop compatibility for the
old code path in new uploader clients.

Test Plan:
Unit tests updated; verified that the `list` and `export` subcommands
still work.

wchargin-branch: uploader-require-experiment-metadata

Author: wchargin

tensorboard/uploader/BUILD

@@ -20,6 +20,7 @@ py_library(
         "//tensorboard:expect_grpc_installed",
         "//tensorboard/uploader/proto:protos_all_py_pb2",
         "//tensorboard/util:grpc_util",
+        "//tensorboard/util:tb_logging",
         "@org_pythonhosted_six",
     ],
 )

tensorboard/uploader/exporter.py

@@ -32,6 +32,7 @@
 from tensorboard.uploader.proto import export_service_pb2
 from tensorboard.uploader import util
 from tensorboard.util import grpc_util
+from tensorboard.util import tb_logging
 
 # Characters that are assumed to be safe in filenames. Note that the
 # server's experiment IDs are base64 encodings of 16-byte blobs, so they
@@ -47,6 +48,8 @@
 # Output filename for scalar data within an experiment directory.
 _FILENAME_SCALARS = "scalars.json"
 
+logger = tb_logging.get_logger()
+
 
 class TensorBoardExporter(object):
     """Exports all of the user's experiment data from TensorBoard.dev.
@@ -115,7 +118,8 @@ def export(self, read_time=None):
         """
         if read_time is None:
             read_time = time.time()
-        for experiment_id in self._request_experiment_ids(read_time):
+        for experiment in list_experiments(self._api, read_time=read_time):
+            experiment_id = experiment.experiment_id
             experiment_dir = _experiment_directory(self._outdir, experiment_id)
             os.mkdir(experiment_dir)
 
@@ -134,18 +138,6 @@ def export(self, read_time=None):
                 else:
                     raise
 
-    def _request_experiment_ids(self, read_time):
-        """Yields all of the calling user's experiment IDs, as strings."""
-        for experiment in list_experiments(self._api, read_time=read_time):
-            if isinstance(experiment, experiment_pb2.Experiment):
-                yield experiment.experiment_id
-            elif isinstance(experiment, six.string_types):
-                yield experiment
-            else:
-                raise AssertionError(
-                    "Unexpected experiment type: %r" % (experiment,)
-                )
-
     def _request_scalar_data(self, experiment_id, read_time):
         """Yields JSON-serializable blocks of scalar data."""
         request = export_service_pb2.StreamExperimentDataRequest()
@@ -191,7 +183,11 @@ def list_experiments(api_client, fieldmask=None, read_time=None):
 
     Yields:
       For each experiment owned by the user, an `experiment_pb2.Experiment`
-      value, or a simple string experiment ID for older servers.
+      value.
+
+    Raises:
+      RuntimeError: If the server returns experiment IDs but no experiments,
+        as in an old, unsupported version of the protocol.
     """
     if read_time is None:
         read_time = time.time()
@@ -206,10 +202,17 @@ def list_experiments(api_client, fieldmask=None, read_time=None):
         if response.experiments:
             for experiment in response.experiments:
                 yield experiment
+        elif response.experiment_ids:
+            raise RuntimeError(
+                "Server sent experiment_ids without experiments: <%r>"
+                % (list(response.experiment_ids),)
+            )
         else:
-            # Old servers.
-            for experiment_id in response.experiment_ids:
-                yield experiment_id
+            # No data: not technically a problem, but not expected.
+            logger.warn(
+                "StreamExperiments RPC returned response with no experiments: <%r>",
+                response,
+            )
 
 
 class OutputDirectoryExistsError(ValueError):

tensorboard/uploader/exporter_test.py

@@ -43,32 +43,29 @@
 from tensorboard.compat.proto import summary_pb2
 
 
+def _make_experiments_response(eids):
+    """Make a `StreamExperimentsResponse` with experiments with only IDs."""
+    response = export_service_pb2.StreamExperimentsResponse()
+    for eid in eids:
+        response.experiments.add(experiment_id=eid)
+    return response
+
+
 class TensorBoardExporterTest(tb_test.TestCase):
     def _create_mock_api_client(self):
         return _create_mock_api_client()
 
-    def _make_experiments_response(self, eids):
-        return export_service_pb2.StreamExperimentsResponse(experiment_ids=eids)
-
     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"]
-                ),
-            ]
+            [_make_experiments_response(["789"])]
         )
 
         def stream_experiments(request, **kwargs):
             del request  # unused
             self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
-            yield export_service_pb2.StreamExperimentsResponse(
-                experiment_ids=["123", "456"]
-            )
-            yield export_service_pb2.StreamExperimentsResponse(
-                experiment_ids=["789"]
-            )
+            yield _make_experiments_response(["123", "456"])
+            yield _make_experiments_response(["789"])
 
         def stream_experiment_data(request, **kwargs):
             self.assertEqual(kwargs["metadata"], grpc_util.version_metadata())
@@ -200,9 +197,7 @@ def test_rejects_dangerous_experiment_ids(self):
 
         def stream_experiments(request, **kwargs):
             del request  # unused
-            yield export_service_pb2.StreamExperimentsResponse(
-                experiment_ids=["../authorized_keys"]
-            )
+            yield _make_experiments_response(["../authorized_keys"])
 
         mock_api_client.StreamExperiments = stream_experiments
 
@@ -229,9 +224,7 @@ def test_fails_nicely_on_stream_experiment_data_timeout(self):
 
         def stream_experiments(request, **kwargs):
             del request  # unused
-            yield export_service_pb2.StreamExperimentsResponse(
-                experiment_ids=[experiment_id]
-            )
+            yield _make_experiments_response([experiment_id])
 
         def stream_experiment_data(request, **kwargs):
             raise test_util.grpc_error(
@@ -260,9 +253,7 @@ def test_stream_experiment_data_passes_through_unexpected_exception(self):
 
         def stream_experiments(request, **kwargs):
             del request  # unused
-            yield export_service_pb2.StreamExperimentsResponse(
-                experiment_ids=[experiment_id]
-            )
+            yield _make_experiments_response([experiment_id])
 
         def stream_experiment_data(request, **kwargs):
             del request  # unused
@@ -288,11 +279,7 @@ def test_handles_outdir_with_no_slash(self):
             os.chdir(self.get_temp_dir())
             mock_api_client = self._create_mock_api_client()
             mock_api_client.StreamExperiments.return_value = iter(
-                [
-                    export_service_pb2.StreamExperimentsResponse(
-                        experiment_ids=["123"]
-                    ),
-                ]
+                [_make_experiments_response(["123"])]
             )
             mock_api_client.StreamExperimentData.return_value = iter(
                 [export_service_pb2.StreamExperimentDataResponse()]
@@ -335,6 +322,7 @@ def test_propagates_mkdir_errors(self):
 
 class ListExperimentsTest(tb_test.TestCase):
     def test_experiment_ids_only(self):
+        # Legacy server behavior; should raise an error.
         mock_api_client = _create_mock_api_client()
 
         def stream_experiments(request, **kwargs):
@@ -347,45 +335,48 @@ def stream_experiments(request, **kwargs):
             )
 
         mock_api_client.StreamExperiments = mock.Mock(wraps=stream_experiments)
-        gen = exporter_lib.list_experiments(mock_api_client)
-        mock_api_client.StreamExperiments.assert_not_called()
-        self.assertEqual(list(gen), ["123", "456", "789"])
+        with self.assertRaises(RuntimeError) as cm:
+            list(exporter_lib.list_experiments(mock_api_client))
+        self.assertIn(repr(["123", "456"]), str(cm.exception))
 
     def test_mixed_experiments_and_ids(self):
         mock_api_client = _create_mock_api_client()
 
         def stream_experiments(request, **kwargs):
             del request  # unused
 
-            # Should include `experiment_ids` when no `experiments` given.
-            response = export_service_pb2.StreamExperimentsResponse()
-            response.experiment_ids.append("123")
-            response.experiment_ids.append("456")
-            yield response
-
             # Should ignore `experiment_ids` in the presence of `experiments`.
             response = export_service_pb2.StreamExperimentsResponse()
             response.experiment_ids.append("999")  # will be omitted
             response.experiments.add(experiment_id="789")
             response.experiments.add(experiment_id="012")
             yield response
 
-            # Should include `experiments` even when no `experiment_ids` are given.
+        mock_api_client.StreamExperiments = mock.Mock(wraps=stream_experiments)
+        gen = exporter_lib.list_experiments(mock_api_client)
+        mock_api_client.StreamExperiments.assert_not_called()
+        expected = [
+            experiment_pb2.Experiment(experiment_id="789"),
+            experiment_pb2.Experiment(experiment_id="012"),
+        ]
+        self.assertEqual(list(gen), expected)
+
+    def test_experiments_only(self):
+        mock_api_client = _create_mock_api_client()
+
+        def stream_experiments(request, **kwargs):
+            del request  # unused
             response = export_service_pb2.StreamExperimentsResponse()
-            response.experiments.add(experiment_id="345")
-            response.experiments.add(experiment_id="678")
+            response.experiments.add(experiment_id="789", name="one")
+            response.experiments.add(experiment_id="012", description="two")
             yield response
 
         mock_api_client.StreamExperiments = mock.Mock(wraps=stream_experiments)
         gen = exporter_lib.list_experiments(mock_api_client)
         mock_api_client.StreamExperiments.assert_not_called()
         expected = [
-            "123",
-            "456",
-            experiment_pb2.Experiment(experiment_id="789"),
-            experiment_pb2.Experiment(experiment_id="012"),
-            experiment_pb2.Experiment(experiment_id="345"),
-            experiment_pb2.Experiment(experiment_id="678"),
+            experiment_pb2.Experiment(experiment_id="789", name="one"),
+            experiment_pb2.Experiment(experiment_id="012", description="two"),
         ]
         self.assertEqual(list(gen), expected)
 

tensorboard/uploader/uploader_main.py

@@ -515,10 +515,6 @@ def execute(self, server_info, channel):
         count = 0
         for experiment in gen:
             count += 1
-            if not isinstance(experiment, experiment_pb2.Experiment):
-                url = server_info_lib.experiment_url(server_info, experiment)
-                print(url)
-                continue
             experiment_id = experiment.experiment_id
             url = server_info_lib.experiment_url(server_info, experiment_id)
             print(url)