From 0f124e93a1ddead16c0018970f34e45c73d5ed81 Mon Sep 17 00:00:00 2001 From: mkovalski Date: Wed, 1 Dec 2021 15:05:33 -0500 Subject: [PATCH] fix: Import error for cloud_profiler (#869) --- .../training_utils/cloud_profiler/__init__.py | 8 +---- .../cloud_profiler/cloud_profiler_utils.py | 21 +++++++++++ .../cloud_profiler/initializer.py | 10 +++++- .../plugins/tensorflow/tf_profiler.py | 30 ++++++---------- setup.py | 8 +++-- tests/unit/aiplatform/test_cloud_profiler.py | 35 +++++++++++++------ 6 files changed, 72 insertions(+), 40 deletions(-) create mode 100644 google/cloud/aiplatform/training_utils/cloud_profiler/cloud_profiler_utils.py diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py b/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py index f5aa40cc34..1b0c5eb925 100644 --- a/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py +++ b/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py @@ -15,13 +15,7 @@ # limitations under the License. # -try: - import google.cloud.aiplatform.training_utils.cloud_profiler.initializer as initializer -except ImportError as err: - raise ImportError( - "Could not load the cloud profiler. To use the profiler, " - 'install the SDK using "pip install google-cloud-aiplatform[cloud-profiler]"' - ) from err +from google.cloud.aiplatform.training_utils.cloud_profiler import initializer """ Initialize the cloud profiler for tensorflow. diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/cloud_profiler_utils.py b/google/cloud/aiplatform/training_utils/cloud_profiler/cloud_profiler_utils.py new file mode 100644 index 0000000000..f7f6e8d8f6 --- /dev/null +++ b/google/cloud/aiplatform/training_utils/cloud_profiler/cloud_profiler_utils.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import_error_msg = ( + "Could not load the cloud profiler. To use the profiler, " + "install the SDK using 'pip install google-cloud-aiplatform[cloud-profiler]'" +) diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/initializer.py b/google/cloud/aiplatform/training_utils/cloud_profiler/initializer.py index 1a098dd2a5..5bb5d19391 100644 --- a/google/cloud/aiplatform/training_utils/cloud_profiler/initializer.py +++ b/google/cloud/aiplatform/training_utils/cloud_profiler/initializer.py @@ -18,7 +18,14 @@ import logging import threading from typing import Optional, Type -from werkzeug import serving + +from google.cloud.aiplatform.training_utils.cloud_profiler import cloud_profiler_utils + +try: + from werkzeug import serving +except ImportError as err: + raise ImportError(cloud_profiler_utils.import_error_msg) from err + from google.cloud.aiplatform.training_utils import environment_variables from google.cloud.aiplatform.training_utils.cloud_profiler import webserver @@ -27,6 +34,7 @@ tf_profiler, ) + # Mapping of available plugins to use _AVAILABLE_PLUGINS = {"tensorflow": tf_profiler.TFProfiler} diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tf_profiler.py b/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tf_profiler.py index 514ae19368..1f0bbccc3b 100644 --- a/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tf_profiler.py +++ b/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tf_profiler.py @@ -17,14 +17,23 @@ """A plugin to handle remote tensoflow profiler sessions for Vertex AI.""" +from google.cloud.aiplatform.training_utils.cloud_profiler import cloud_profiler_utils + +try: + import tensorflow as tf + from tensorboard_plugin_profile.profile_plugin import ProfilePlugin +except ImportError as err: + raise ImportError(cloud_profiler_utils.import_error_msg) from err + import argparse from collections import namedtuple import importlib.util import json import logging -import tensorboard.plugins.base_plugin as tensorboard_base_plugin from typing import Callable, Dict, Optional from urllib import parse + +import tensorboard.plugins.base_plugin as tensorboard_base_plugin from werkzeug import Response from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader @@ -54,8 +63,6 @@ def _get_tf_versioning() -> Optional[Version]: Returns: A version object if finding the version was successful, None otherwise. """ - import tensorflow as tf - version = tf.__version__ versioning = version.split(".") @@ -269,8 +276,6 @@ class TFProfiler(base_plugin.BasePlugin): def __init__(self): """Build a TFProfiler object.""" - from tensorboard_plugin_profile.profile_plugin import ProfilePlugin - context = _create_profiling_context() self._profile_request_sender: profile_uploader.ProfileRequestSender = tensorboard_api.create_profile_request_sender() self._profile_plugin: ProfilePlugin = ProfilePlugin(context) @@ -317,20 +322,7 @@ def capture_profile_wrapper( @staticmethod def setup() -> None: - """Sets up the plugin. - - Raises: - ImportError: Tensorflow could not be imported. - """ - try: - import tensorflow as tf - except ImportError as err: - raise ImportError( - "Could not import tensorflow for profile usage. " - "To use profiler, install the SDK using " - '"pip install google-cloud-aiplatform[cloud_profiler]"' - ) from err - + """Sets up the plugin.""" tf.profiler.experimental.server.start( int(environment_variables.tf_profiler_port) ) diff --git a/setup.py b/setup.py index d0daa259f3..4ef6968114 100644 --- a/setup.py +++ b/setup.py @@ -36,7 +36,11 @@ tensorboard_extra_require = ["tensorflow >=2.3.0, <=2.5.0"] metadata_extra_require = ["pandas >= 1.0.0"] xai_extra_require = ["tensorflow >=2.3.0, <=2.5.0"] -profiler_extra_require = ["tensorboard-plugin-profile", "tensorflow >=2.4.0"] +profiler_extra_require = [ + "tensorboard-plugin-profile >= 2.4.0", + "werkzeug >= 2.0.0", + "tensorflow >=2.4.0", +] full_extra_require = list( set(tensorboard_extra_require + metadata_extra_require + xai_extra_require) @@ -84,7 +88,7 @@ "tensorboard": tensorboard_extra_require, "testing": testing_extra_require, "xai": xai_extra_require, - "cloud_profiler": profiler_extra_require, + "cloud-profiler": profiler_extra_require, }, python_requires=">=3.6", scripts=[], diff --git a/tests/unit/aiplatform/test_cloud_profiler.py b/tests/unit/aiplatform/test_cloud_profiler.py index e540279bf9..3cde7a1296 100644 --- a/tests/unit/aiplatform/test_cloud_profiler.py +++ b/tests/unit/aiplatform/test_cloud_profiler.py @@ -15,9 +15,9 @@ # limitations under the License. # -from importlib import reload import importlib.util import json +import sys import threading from typing import List, Optional @@ -75,6 +75,10 @@ def _create_mock_plugin( return mock_plugin +def _find_child_modules(root_module): + return [module for module in sys.modules.keys() if module.startswith(root_module)] + + @pytest.fixture def tf_profile_plugin_mock(): """Mock the tensorboard profile plugin""" @@ -203,10 +207,6 @@ def testSetup(self): assert server_mock.call_count == 1 - def testSetupRaiseImportError(self): - with mock.patch.dict("sys.modules", {"tensorflow": None}): - self.assertRaises(ImportError, TFProfiler.setup) - def testPostSetupChecksFail(self): tf_profiler.environment_variables.cluster_spec = {} assert not TFProfiler.post_setup_check() @@ -359,13 +359,26 @@ def start_response(status, headers): # Initializer tests class TestInitializer(unittest.TestCase): - # Tests for building the plugin - def test_init_failed_import(self): - with mock.patch.dict( - "sys.modules", - {"google.cloud.aiplatform.training_utils.cloud_profiler.initializer": None}, + def testImportError(self): + # Unloads any of the cloud profiler sub-modules + for mod in _find_child_modules( + "google.cloud.aiplatform.training_utils.cloud_profiler" ): - self.assertRaises(ImportError, reload, training_utils.cloud_profiler) + del sys.modules[mod] + + # Modules to be mocked out + for mock_module in [ + "tensorflow", + "tensorboard_plugin_profile.profile_plugin", + "werkzeug", + ]: + with self.subTest(): + with mock.patch.dict("sys.modules", {mock_module: None}): + with self.assertRaises(ImportError) as cm: + importlib.import_module( + "google.cloud.aiplatform.training_utils.cloud_profiler" + ) + assert "Could not load the cloud profiler" in cm.exception.msg def test_build_plugin_fail_initialize(self): plugin = _create_mock_plugin()