uploader: add metadata to experiment listing RPC (#2906)

Summary:
This extends the `StreamExperiments` RPC such that the client can
specify a set of additional metadata fields that the server should
include, like “creation time” or “number of scalar points”.

The format is both forward- and backward-compatible. Servers are
expected to send responses with both `experiment_ids` and `experiments`
until we drop support for clients that do not support `experiments`, at
which point they need only send `experiments`.

Test Plan:
Unit test added to simulate the future behavior of servers.

wchargin-branch: streamexperiments-metadata

Author: wchargin

tensorboard/uploader/exporter.py

@@ -177,8 +177,11 @@ def list_experiments(api_client, read_time=None):
   stream = api_client.StreamExperiments(
       request, metadata=grpc_util.version_metadata())
   for response in stream:
-    for experiment_id in response.experiment_ids:
-      yield experiment_id
+    if not response.experiments:
+      for experiment_id in response.experiment_ids:
+        yield experiment_id
+    for experiment in response.experiments:
+      yield experiment.experiment_id
 
 
 class OutputDirectoryExistsError(ValueError):

tensorboard/uploader/exporter_test.py

@@ -317,7 +317,7 @@ def test_propagates_mkdir_errors(self):
 
 class ListExperimentsTest(tb_test.TestCase):
 
-  def test(self):
+  def test_experiment_ids_only(self):
     mock_api_client = _create_mock_api_client()
 
     def stream_experiments(request, **kwargs):
@@ -332,6 +332,36 @@ def stream_experiments(request, **kwargs):
     mock_api_client.StreamExperiments.assert_not_called()
     self.assertEqual(list(gen), ["123", "456", "789"])
 
+  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.
+      response = export_service_pb2.StreamExperimentsResponse()
+      response.experiments.add(experiment_id="345")
+      response.experiments.add(experiment_id="678")
+      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()
+    self.assertEqual(list(gen), ["123", "456", "789", "012", "345", "678"])
+
 
 class MkdirPTest(tb_test.TestCase):
 

tensorboard/uploader/proto/export_service.proto

@@ -27,19 +27,80 @@ message StreamExperimentsRequest {
   string user_id = 2;
   // Limits the number of experiment IDs returned. This is useful to check if
   // user might have any data by setting limit=1. Also useful to preview the
-  // list of experiments.
+  // list of experiments. TODO(@karthikv2k): Support pagination.
   int64 limit = 3;
-  // TODO(@karthikv2k): Support pagination.
+  // Field mask for what experiment data to return via the `experiments` field
+  // on the response. If not specified, this should be interpreted the same as
+  // an empty message: i.e., only the experiment ID should be returned.
+  ExperimentMask experiments_mask = 4;
 }
 
-// Streams experiment IDs returned from TensorBoard.dev.
+// Streams experiment metadata (ID, creation time, etc.) from TensorBoard.dev.
 message StreamExperimentsResponse {
-  // List of experiment IDs for the experiments owned by the user. The entire
-  // list of experiments owned by the user is streamed in batches and each batch
-  // contains a list of experiment IDs. A consumer of this stream needs to
-  // concatenate all these lists to get the full response. The order of
-  // experiment IDs in the stream is not defined.
+  // Deprecated in favor of `experiments`. If a response has `experiments` set,
+  // clients should ignore `experiment_ids` entirely. Otherwise, clients should
+  // treat `experiment_ids` as a list of `experiments` for which only the
+  // `experiment_id` field is set, with the understanding that the other fields
+  // were not populated regardless of the requested field mask.
+  //
+  // For example, the following responses should be treated the same:
+  //
+  //     # Response 1
+  //     experiment_ids: "123"
+  //     experiment_ids: "456"
+  //
+  //     # Response 2
+  //     experiments { experiment_id: "123" }
+  //     experiments { experiment_id: "456" }
+  //
+  //     # Response 3
+  //     experiment_ids: "789"
+  //     experiments { experiment_id: "123" }
+  //     experiments { experiment_id: "456" }
+  //
+  // See documentation on `experiments` for batching semantics.
   repeated string experiment_ids = 1;
+  // List of experiments owned by the user. The entire list of experiments
+  // owned by the user is streamed in batches and each batch contains a list of
+  // experiments. A consumer of this stream needs to concatenate all these
+  // lists to get the full response. The order of experiments in the stream is
+  // not defined. Every response will contain at least one experiment.
+  //
+  // These messages may be partially populated, in accordance with the field
+  // mask given in the request.
+  repeated Experiment experiments = 2;
+}
+
+// Metadata about an experiment.
+message Experiment {
+  // Permanent ID of this experiment; e.g.: "AdYd1TgeTlaLWXx6I8JUbA".
+  string experiment_id = 1;
+  // The time that the experiment was created.
+  google.protobuf.Timestamp create_time = 2;
+  // The time that the experiment was last modified: i.e., the most recent time
+  // that scalars were added to the experiment.
+  google.protobuf.Timestamp update_time = 3;
+  // The number of scalars in this experiment, across all time series.
+  int64 num_scalars = 4;
+  // The number of distinct run names in this experiment.
+  int64 num_runs = 5;
+  // The number of distinct tag names in this experiment. A tag name that
+  // appears in multiple runs will be counted only once.
+  int64 num_tags = 6;
+}
+
+// Field mask for `Experiment`. The `experiment_id` field is always implicitly
+// considered to be requested. Other fields of `Experiment` will be populated
+// if their corresponding bits in the `ExperimentMask` are set. The server may
+// choose to populate fields that are not explicitly requested.
+message ExperimentMask {
+  reserved 1;
+  reserved "experiment_id";
+  bool create_time = 2;
+  bool update_time = 3;
+  bool num_scalars = 4;
+  bool num_runs = 5;
+  bool num_tags = 6;
 }
 
 // Request to stream scalars from all the runs and tags in an experiment.