|
1 | 1 | # -*- coding: utf-8 -*-
|
2 | 2 |
|
3 |
| -# Copyright 2021 Google LLC |
| 3 | +# Copyright 2023 Google LLC |
4 | 4 | #
|
5 | 5 | # Licensed under the Apache License, Version 2.0 (the "License");
|
6 | 6 | # you may not use this file except in compliance with the License.
|
|
15 | 15 | # limitations under the License.
|
16 | 16 | #
|
17 | 17 | """Uploads a TensorBoard logdir to TensorBoard.gcp."""
|
| 18 | + |
18 | 19 | import abc
|
19 | 20 | from collections import defaultdict
|
20 | 21 | import functools
|
21 | 22 | import logging
|
22 | 23 | import os
|
23 |
| -import time |
24 | 24 | import re
|
25 |
| -from typing import ( |
26 |
| - Dict, |
27 |
| - FrozenSet, |
28 |
| - Generator, |
29 |
| - Iterable, |
30 |
| - Optional, |
31 |
| - ContextManager, |
32 |
| - Tuple, |
33 |
| -) |
| 25 | +import time |
| 26 | +from typing import ContextManager, Dict, FrozenSet, Generator, Iterable, Optional, Tuple |
34 | 27 | import uuid
|
35 | 28 |
|
| 29 | +from google.api_core import exceptions |
| 30 | +from google.cloud import storage |
| 31 | +from google.cloud.aiplatform import base |
| 32 | +from google.cloud.aiplatform.compat.services import ( |
| 33 | + tensorboard_service_client, |
| 34 | +) |
| 35 | +from google.cloud.aiplatform.compat.types import tensorboard_data |
| 36 | +from google.cloud.aiplatform.compat.types import tensorboard_experiment |
| 37 | +from google.cloud.aiplatform.compat.types import tensorboard_service |
| 38 | +from google.cloud.aiplatform.compat.types import tensorboard_time_series |
| 39 | +from google.cloud.aiplatform.tensorboard import uploader_utils |
| 40 | +from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import ( |
| 41 | + profile_uploader, |
| 42 | +) |
36 | 43 | import grpc
|
| 44 | +import tensorflow as tf |
| 45 | + |
| 46 | +from google.protobuf import timestamp_pb2 as timestamp |
| 47 | +from google.protobuf import message |
37 | 48 | from tensorboard.backend import process_graph
|
38 | 49 | from tensorboard.backend.event_processing.plugin_event_accumulator import (
|
39 | 50 | directory_loader,
|
40 | 51 | )
|
41 | 52 | from tensorboard.backend.event_processing.plugin_event_accumulator import (
|
42 | 53 | event_file_loader,
|
43 | 54 | )
|
44 |
| -from tensorboard.backend.event_processing.plugin_event_accumulator import io_wrapper |
| 55 | +from tensorboard.backend.event_processing.plugin_event_accumulator import ( |
| 56 | + io_wrapper, |
| 57 | +) |
45 | 58 | from tensorboard.compat.proto import graph_pb2
|
46 | 59 | from tensorboard.compat.proto import summary_pb2
|
47 | 60 | from tensorboard.compat.proto import types_pb2
|
|
52 | 65 | from tensorboard.uploader.proto import server_info_pb2
|
53 | 66 | from tensorboard.util import tb_logging
|
54 | 67 | from tensorboard.util import tensor_util
|
55 |
| -import tensorflow as tf |
56 | 68 |
|
57 |
| -from google.api_core import exceptions |
58 |
| -from google.cloud import storage |
59 |
| -from google.cloud.aiplatform.compat.services import tensorboard_service_client |
60 |
| -from google.cloud.aiplatform.compat.types import tensorboard_data |
61 |
| -from google.cloud.aiplatform.compat.types import tensorboard_experiment |
62 |
| -from google.cloud.aiplatform.compat.types import tensorboard_service |
63 |
| -from google.cloud.aiplatform.compat.types import tensorboard_time_series |
64 |
| -from google.cloud.aiplatform.tensorboard import uploader_utils |
65 |
| -from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader |
66 |
| -from google.protobuf import message |
67 |
| -from google.protobuf import timestamp_pb2 as timestamp |
| 69 | +_LOGGER = base.Logger(__name__) |
68 | 70 |
|
69 | 71 | TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
|
70 | 72 |
|
@@ -189,6 +191,7 @@ def __init__(
|
189 | 191 | self._allowed_plugins = frozenset(allowed_plugins)
|
190 | 192 | self._run_name_prefix = run_name_prefix
|
191 | 193 | self._is_brand_new_experiment = False
|
| 194 | + self._continue_uploading = True |
192 | 195 |
|
193 | 196 | self._upload_limits = upload_limits
|
194 | 197 | if not self._upload_limits:
|
@@ -388,20 +391,22 @@ def start_uploading(self):
|
388 | 391 | "performance."
|
389 | 392 | )
|
390 | 393 |
|
391 |
| - while True: |
| 394 | + while self._continue_uploading: |
392 | 395 | self._logdir_poll_rate_limiter.tick()
|
393 | 396 | self._upload_once()
|
394 | 397 | if self._one_shot:
|
395 | 398 | break
|
396 | 399 | if self._one_shot and not self._tracker.has_data():
|
397 | 400 | logger.warning(
|
398 |
| - "One-shot mode was used on a logdir (%s) " |
399 |
| - "without any uploadable data" % self._logdir |
| 401 | + "One-shot mode was used on a logdir (%s) without any uploadable data" |
| 402 | + % self._logdir |
400 | 403 | )
|
401 | 404 |
|
| 405 | + def _end_uploading(self): |
| 406 | + self._continue_uploading = False |
| 407 | + |
402 | 408 | def _pre_create_runs_and_time_series(self):
|
403 |
| - """ |
404 |
| - Iterates though the log dir to collect TensorboardRuns and |
| 409 | + """Iterates though the log dir to collect TensorboardRuns and |
405 | 410 | TensorboardTimeSeries that need to be created, and creates them in batch
|
406 | 411 | to speed up uploading later on.
|
407 | 412 | """
|
|
0 commit comments