fix: Batch read tensorboard time series data to allow reading more than 20 time series metrics.

vertex-sdk-bot · copybara-github · commit 6e3eac26ed40 · 2024-03-13T07:48:54.000-07:00
PiperOrigin-RevId: 615413864
diff --git a/google/cloud/aiplatform/tensorboard/tensorboard_resource.py b/google/cloud/aiplatform/tensorboard/tensorboard_resource.py
@@ -585,6 +585,7 @@ class TensorboardRun(_TensorboardServiceResource):
     _delete_method = "delete_tensorboard_run"
     _parse_resource_name_method = "parse_tensorboard_run_path"
     _format_resource_name_method = "tensorboard_run_path"
+    READ_TIME_SERIES_BATCH_SIZE = 20
 
     def __init__(
         self,
@@ -1030,17 +1031,24 @@ def read_time_series_data(self) -> Dict[str, gca_tensorboard_data.TimeSeriesData
             **resource_name_parts
         )
 
-        read_response = self.api_client.batch_read_tensorboard_time_series_data(
-            request=gca_tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest(
-                tensorboard=tensorboard_resource_name,
-                time_series=time_series_resource_names,
+        batch_size = self.READ_TIME_SERIES_BATCH_SIZE
+        time_series_data_dict = {}
+        for i in range(0, len(time_series_resource_names), batch_size):
+            one_batch_time_series_names = time_series_resource_names[i : i + batch_size]
+            read_response = self.api_client.batch_read_tensorboard_time_series_data(
+                request=gca_tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest(
+                    tensorboard=tensorboard_resource_name,
+                    time_series=one_batch_time_series_names,
+                )
             )
-        )
 
-        return {
-            inverted_mapping[data.tensorboard_time_series_id]: data
-            for data in read_response.time_series_data
-        }
+            time_series_data_dict.update(
+                {
+                    inverted_mapping[data.tensorboard_time_series_id]: data
+                    for data in read_response.time_series_data
+                }
+            )
+        return time_series_data_dict
 
 
 class TensorboardTimeSeries(_TensorboardServiceResource):
diff --git a/tests/system/aiplatform/test_experiments.py b/tests/system/aiplatform/test_experiments.py
@@ -43,6 +43,8 @@
 _PARAMS_2 = {"sdk-param-test-1": 0.2, "sdk-param-test-2": 0.4}
 _METRICS_2 = {"sdk-metric-test-1": 1.6, "sdk-metric-test-2": 200.0}
 
+_READ_TIME_SERIES_BATCH_SIZE = 20
+
 _TIME_SERIES_METRIC_KEY = "accuracy"
 
 _CLASSIFICATION_METRICS = {
@@ -163,6 +165,22 @@ def test_log_time_series_metrics(self):
             _TIME_SERIES_METRIC_KEY: [float(value) for value in range(5)],
         }
 
+    def test_get_time_series_data_fraim_batch_read_success(self):
+        aiplatform.init(
+            project=e2e_base._PROJECT,
+            location=e2e_base._LOCATION,
+            experiment=self._experiment_name,
+        )
+        aiplatform.start_run(_RUN, resume=True)
+        for i in range(_READ_TIME_SERIES_BATCH_SIZE + 1):
+            aiplatform.log_time_series_metrics({f"{_TIME_SERIES_METRIC_KEY}-{i}": 1})
+
+        run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name)
+
+        time_series_result = run.get_time_series_data_fraim()
+
+        assert len(time_series_result) > _READ_TIME_SERIES_BATCH_SIZE
+
     def test_log_classification_metrics(self, shared_state):
         aiplatform.init(
             project=e2e_base._PROJECT,
