Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: updating Tensorboard related code to use v1 #851

Merged
merged 1 commit into from Nov 19, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
1 change: 1 addition & 0 deletions google/cloud/aiplatform/compat/__init__.py
Expand Up @@ -95,6 +95,7 @@
services.pipeline_service_client = services.pipeline_service_client_v1
services.prediction_service_client = services.prediction_service_client_v1
services.specialist_pool_service_client = services.specialist_pool_service_client_v1
services.tensorboard_service_client = services.tensorboard_service_client_v1

types.accelerator_type = types.accelerator_type_v1
types.annotation = types.annotation_v1
Expand Down
Expand Up @@ -38,20 +38,14 @@
import tensorflow as tf

from google.cloud import storage
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
from google.cloud.aiplatform.compat.types import (
tensorboard_data_v1beta1 as tensorboard_data,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_service_v1beta1 as tensorboard_service,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_time_series_v1beta1 as tensorboard_time_series,
)
from google.cloud.aiplatform.compat.services import tensorboard_service_client
from google.cloud.aiplatform.compat.types import tensorboard_data
from google.cloud.aiplatform.compat.types import tensorboard_service
from google.cloud.aiplatform.compat.types import tensorboard_time_series
from google.cloud.aiplatform.tensorboard import uploader_utils
from google.protobuf import timestamp_pb2 as timestamp

TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient
TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient

logger = tb_logging.get_logger()

Expand Down
20 changes: 6 additions & 14 deletions google/cloud/aiplatform/tensorboard/uploader.py
Expand Up @@ -56,25 +56,17 @@

from google.api_core import exceptions
from google.cloud import storage
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
from google.cloud.aiplatform.compat.types import (
tensorboard_data_v1beta1 as tensorboard_data,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_experiment_v1beta1 as tensorboard_experiment,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_service_v1beta1 as tensorboard_service,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_time_series_v1beta1 as tensorboard_time_series,
)
from google.cloud.aiplatform.compat.services import tensorboard_service_client
from google.cloud.aiplatform.compat.types import tensorboard_data
from google.cloud.aiplatform.compat.types import tensorboard_experiment
from google.cloud.aiplatform.compat.types import tensorboard_service
from google.cloud.aiplatform.compat.types import tensorboard_time_series
from google.cloud.aiplatform.tensorboard import uploader_utils
from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
from google.protobuf import message
from google.protobuf import timestamp_pb2 as timestamp

TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient
TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient

# Minimum length of a logdir polling cycle in seconds. Shorter cycles will
# sleep to avoid spinning over the logdir, which isn't great for disks and can
Expand Down
20 changes: 7 additions & 13 deletions google/cloud/aiplatform/tensorboard/uploader_utils.py
Expand Up @@ -29,18 +29,12 @@

from google.api_core import exceptions
from google.cloud import storage
from google.cloud.aiplatform.compat.types import (
tensorboard_run_v1beta1 as tensorboard_run,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_service_v1beta1 as tensorboard_service,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_time_series_v1beta1 as tensorboard_time_series,
)
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1

TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient
from google.cloud.aiplatform.compat.types import tensorboard_run
from google.cloud.aiplatform.compat.types import tensorboard_service
from google.cloud.aiplatform.compat.types import tensorboard_time_series
from google.cloud.aiplatform.compat.services import tensorboard_service_client

TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient

logger = tb_logging.get_logger()
logger.setLevel(logging.WARNING)
Expand Down Expand Up @@ -201,7 +195,7 @@ def _create_or_get_run_resource(
Required. The display name of this run.

Returns:
tb_run (google.cloud.aiplatform_v1beta1.types.TensorboardRun):
tb_run (tensorboard_run.TensorboardRun):
The TensorboardRun given the run_name.

Raises:
Expand Down
26 changes: 12 additions & 14 deletions tests/unit/aiplatform/test_uploader.py
Expand Up @@ -46,22 +46,20 @@
from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
import google.cloud.aiplatform.tensorboard.uploader as uploader_lib
from google.cloud import storage
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
from google.cloud.aiplatform_v1beta1.services.tensorboard_service.transports import (
grpc as transports_grpc,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_data_v1beta1 as tensorboard_data,
tensorboard_service_v1beta1 as tensorboard_service,
from google.cloud.aiplatform_v1.services.tensorboard_service import (
client as tensorboard_service_client,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_experiment_v1beta1 as tensorboard_experiment_type,
from google.cloud.aiplatform_v1.services.tensorboard_service.transports import (
grpc as transports_grpc,
)
from google.cloud.aiplatform.compat.types import (
tensorboard_run_v1beta1 as tensorboard_run_type,
from google.cloud.aiplatform_v1.types import tensorboard_data
from google.cloud.aiplatform_v1.types import tensorboard_service
from google.cloud.aiplatform_v1.types import (
tensorboard_experiment as tensorboard_experiment_type,
yfang1 marked this conversation as resolved.
Show resolved Hide resolved
)
from google.cloud.aiplatform.compat.types import (
tensorboard_time_series_v1beta1 as tensorboard_time_series_type,
from google.cloud.aiplatform_v1.types import tensorboard_run as tensorboard_run_type
from google.cloud.aiplatform_v1.types import (
tensorboard_time_series as tensorboard_time_series_type,
yfang1 marked this conversation as resolved.
Show resolved Hide resolved
)
from google.protobuf import timestamp_pb2
from google.protobuf import message
Expand Down Expand Up @@ -160,7 +158,7 @@ def create_tensorboard_time_series(
service_descriptors=[], time=grpc_testing.strict_real_time()
)
mock_client = mock.Mock(
spec=tensorboard_service_client_v1beta1.TensorboardServiceClient(
spec=tensorboard_service_client.TensorboardServiceClient(
transport=transports_grpc.TensorboardServiceGrpcTransport(
channel=test_channel
)
Expand Down