From 1368556c4dcacbb129b8484bc0867a09eaf4b33f Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Mon, 29 Nov 2021 20:21:49 +0000 Subject: [PATCH 1/9] fix import error involving absolute imports: bpo-30024 --- .../cloud/aiplatform/training_utils/cloud_profiler/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py b/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py index f5aa40cc34..806252fda9 100644 --- a/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py +++ b/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py @@ -16,7 +16,7 @@ # try: - import google.cloud.aiplatform.training_utils.cloud_profiler.initializer as initializer + import google.cloud.aiplatform.training_utils.cloud_profiler.initializer except ImportError as err: raise ImportError( "Could not load the cloud profiler. To use the profiler, " From 7d8ef46950ac348aa32ff54e47532a033988141f Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Mon, 29 Nov 2021 20:21:49 +0000 Subject: [PATCH 2/9] fix: import error involving absolute imports: fixes #867 --- .../cloud/aiplatform/training_utils/cloud_profiler/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py b/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py index f5aa40cc34..806252fda9 100644 --- a/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py +++ b/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py @@ -16,7 +16,7 @@ # try: - import google.cloud.aiplatform.training_utils.cloud_profiler.initializer as initializer + import google.cloud.aiplatform.training_utils.cloud_profiler.initializer except ImportError as err: raise ImportError( "Could not load the cloud profiler. To use the profiler, " From d76cef1385078c44c12ed1ad9d5765a4322c949f Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Mon, 29 Nov 2021 23:55:15 +0000 Subject: [PATCH 3/9] fix: Move the import checks to tf profiler module --- .../training_utils/cloud_profiler/__init__.py | 8 +------ .../plugins/tensorflow/tf_profiler.py | 10 +++++++-- tests/unit/aiplatform/test_cloud_profiler.py | 21 ++++++++++++------- 3 files changed, 22 insertions(+), 17 deletions(-) diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py b/google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py index 806252fda9..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 -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/plugins/tensorflow/tf_profiler.py b/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tf_profiler.py index 514ae19368..565859d29a 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,6 +17,14 @@ """A plugin to handle remote tensoflow profiler sessions for Vertex AI.""" +try: + from tensorboard_plugin_profile.profile_plugin import ProfilePlugin +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 + import argparse from collections import namedtuple import importlib.util @@ -269,8 +277,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) diff --git a/tests/unit/aiplatform/test_cloud_profiler.py b/tests/unit/aiplatform/test_cloud_profiler.py index e540279bf9..9515c42088 100644 --- a/tests/unit/aiplatform/test_cloud_profiler.py +++ b/tests/unit/aiplatform/test_cloud_profiler.py @@ -130,6 +130,19 @@ class TestProfilerPlugin(unittest.TestCase): def setUp(self): setupProfilerEnvVars() + def testImportError(self): + orig_find_spec = importlib.util.find_spec + + def profile_import_mock(name, *args, **kwargs): + if name == "tensorboard_plugin_profile": + return None + return orig_find_spec(name, *args, **kwargs) + + with mock.patch.dict( + "sys.modules", {"tensorboard_plugin_profile.profile_plugin": None} + ): + self.assertRaises(ImportError, reload, tf_profiler) + # Environment variable tests def testCanInitializeProfilerPortUnset(self): tf_profiler.environment_variables.tf_profiler_port = None @@ -359,14 +372,6 @@ 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}, - ): - self.assertRaises(ImportError, reload, training_utils.cloud_profiler) - def test_build_plugin_fail_initialize(self): plugin = _create_mock_plugin() plugin.can_initialize.return_value = False From 91ef8d59d19d7b1a4f1c01434e2987575a19b4e7 Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Tue, 30 Nov 2021 00:12:15 +0000 Subject: [PATCH 4/9] fix: remove unnecessary importlib mock from tf_profiler test --- tests/unit/aiplatform/test_cloud_profiler.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/unit/aiplatform/test_cloud_profiler.py b/tests/unit/aiplatform/test_cloud_profiler.py index 9515c42088..391acc03e0 100644 --- a/tests/unit/aiplatform/test_cloud_profiler.py +++ b/tests/unit/aiplatform/test_cloud_profiler.py @@ -131,13 +131,6 @@ def setUp(self): setupProfilerEnvVars() def testImportError(self): - orig_find_spec = importlib.util.find_spec - - def profile_import_mock(name, *args, **kwargs): - if name == "tensorboard_plugin_profile": - return None - return orig_find_spec(name, *args, **kwargs) - with mock.patch.dict( "sys.modules", {"tensorboard_plugin_profile.profile_plugin": None} ): From 97d41d4af505e4aba6d830eb369494bfa5da80c1 Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Tue, 30 Nov 2021 17:57:00 +0000 Subject: [PATCH 5/9] fix: add werkzeug to profiler requirements and add additional missing module to test --- .../cloud_profiler/plugins/tensorflow/tf_profiler.py | 2 +- setup.py | 6 +++++- tests/unit/aiplatform/test_cloud_profiler.py | 8 ++++---- 3 files changed, 10 insertions(+), 6 deletions(-) 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 565859d29a..fa8d0ab38d 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 @@ -19,6 +19,7 @@ try: from tensorboard_plugin_profile.profile_plugin import ProfilePlugin + from werkzeug import Response except ImportError as err: raise ImportError( "Could not load the cloud profiler. To use the profiler, " @@ -33,7 +34,6 @@ import tensorboard.plugins.base_plugin as tensorboard_base_plugin from typing import Callable, Dict, Optional from urllib import parse -from werkzeug import Response from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader from google.cloud.aiplatform.training_utils import environment_variables diff --git a/setup.py b/setup.py index d0daa259f3..9699381779 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", + "werkzeug", + "tensorflow >=2.4.0", +] full_extra_require = list( set(tensorboard_extra_require + metadata_extra_require + xai_extra_require) diff --git a/tests/unit/aiplatform/test_cloud_profiler.py b/tests/unit/aiplatform/test_cloud_profiler.py index 391acc03e0..9d5c26afa1 100644 --- a/tests/unit/aiplatform/test_cloud_profiler.py +++ b/tests/unit/aiplatform/test_cloud_profiler.py @@ -131,10 +131,10 @@ def setUp(self): setupProfilerEnvVars() def testImportError(self): - with mock.patch.dict( - "sys.modules", {"tensorboard_plugin_profile.profile_plugin": None} - ): - self.assertRaises(ImportError, reload, tf_profiler) + for mock_module in ["tensorboard_plugin_profile.profile_plugin", "werkzeug"]: + with self.subTest(): + with mock.patch.dict("sys.modules", {mock_module: None}): + self.assertRaises(ImportError, reload, tf_profiler) # Environment variable tests def testCanInitializeProfilerPortUnset(self): From b9353a88640532f48dae326cfc08497627c13762 Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Wed, 1 Dec 2021 17:07:15 +0000 Subject: [PATCH 6/9] fix: version requirements on setup.py, catching exception and checking verbiage in error message --- .../cloud_profiler/cloud_profiler_utils.py | 21 ++++++++++++++ .../cloud_profiler/initializer.py | 10 ++++++- .../plugins/tensorflow/tf_profiler.py | 16 +++++----- setup.py | 4 +-- tests/unit/aiplatform/test_cloud_profiler.py | 29 ++++++++++++++----- 5 files changed, 61 insertions(+), 19 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/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 fa8d0ab38d..c983258062 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,15 +17,6 @@ """A plugin to handle remote tensoflow profiler sessions for Vertex AI.""" -try: - from tensorboard_plugin_profile.profile_plugin import ProfilePlugin - from werkzeug import Response -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 - import argparse from collections import namedtuple import importlib.util @@ -37,12 +28,19 @@ from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader from google.cloud.aiplatform.training_utils import environment_variables +from google.cloud.aiplatform.training_utils.cloud_profiler import cloud_profiler_utils from google.cloud.aiplatform.training_utils.cloud_profiler import wsgi_types from google.cloud.aiplatform.training_utils.cloud_profiler.plugins import base_plugin from google.cloud.aiplatform.training_utils.cloud_profiler.plugins.tensorflow import ( tensorboard_api, ) +try: + from tensorboard_plugin_profile.profile_plugin import ProfilePlugin + from werkzeug import Response +except ImportError as err: + raise ImportError(cloud_profiler_utils.import_error_msg) from err + # TF verison information. Version = namedtuple("Version", ["major", "minor", "patch"]) diff --git a/setup.py b/setup.py index 9699381779..2cf78ace6d 100644 --- a/setup.py +++ b/setup.py @@ -37,8 +37,8 @@ metadata_extra_require = ["pandas >= 1.0.0"] xai_extra_require = ["tensorflow >=2.3.0, <=2.5.0"] profiler_extra_require = [ - "tensorboard-plugin-profile", - "werkzeug", + "tensorboard-plugin-profile >= 2.4.0", + "werkzeug >= 2.0.0", "tensorflow >=2.4.0", ] diff --git a/tests/unit/aiplatform/test_cloud_profiler.py b/tests/unit/aiplatform/test_cloud_profiler.py index 9d5c26afa1..e8bdb79fc2 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""" @@ -130,12 +134,6 @@ class TestProfilerPlugin(unittest.TestCase): def setUp(self): setupProfilerEnvVars() - def testImportError(self): - for mock_module in ["tensorboard_plugin_profile.profile_plugin", "werkzeug"]: - with self.subTest(): - with mock.patch.dict("sys.modules", {mock_module: None}): - self.assertRaises(ImportError, reload, tf_profiler) - # Environment variable tests def testCanInitializeProfilerPortUnset(self): tf_profiler.environment_variables.tf_profiler_port = None @@ -365,6 +363,23 @@ def start_response(status, headers): # Initializer tests class TestInitializer(unittest.TestCase): + def testImportError(self): + # Unloads any of the cloud profiler sub-modules + for mod in _find_child_modules( + "google.cloud.aiplatform.training_utils.cloud_profiler" + ): + del sys.modules[mod] + + # Modules to be mocked out + for mock_module in ["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() plugin.can_initialize.return_value = False From 54e38394151bba1041bdaf58057616b7398adb5f Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Wed, 1 Dec 2021 17:17:13 +0000 Subject: [PATCH 7/9] setup using cloud-profiler rather than cloud_profiler --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 2cf78ace6d..4ef6968114 100644 --- a/setup.py +++ b/setup.py @@ -88,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=[], From 71a2c06073b54dd775052cae35ff1abe8984e577 Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Wed, 1 Dec 2021 17:54:14 +0000 Subject: [PATCH 8/9] fix: Add tensorflow to module check --- .../plugins/tensorflow/tf_profiler.py | 18 ++---------------- tests/unit/aiplatform/test_cloud_profiler.py | 10 +++++----- 2 files changed, 7 insertions(+), 21 deletions(-) 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 c983258062..2d9536c947 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 @@ -38,6 +38,7 @@ try: from tensorboard_plugin_profile.profile_plugin import ProfilePlugin from werkzeug import Response + import tensorflow as tf except ImportError as err: raise ImportError(cloud_profiler_utils.import_error_msg) from err @@ -60,8 +61,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(".") @@ -321,20 +320,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/tests/unit/aiplatform/test_cloud_profiler.py b/tests/unit/aiplatform/test_cloud_profiler.py index e8bdb79fc2..3cde7a1296 100644 --- a/tests/unit/aiplatform/test_cloud_profiler.py +++ b/tests/unit/aiplatform/test_cloud_profiler.py @@ -207,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() @@ -371,7 +367,11 @@ def testImportError(self): del sys.modules[mod] # Modules to be mocked out - for mock_module in ["tensorboard_plugin_profile.profile_plugin", "werkzeug"]: + 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: From c3c2c9e1c035a608572748dd7d0b6f811cfffd54 Mon Sep 17 00:00:00 2001 From: Michael Kovalski Date: Wed, 1 Dec 2021 19:26:12 +0000 Subject: [PATCH 9/9] fix: Adding tensorflow and tensorboard as top level import for catching cloud-profiler import errors --- .../plugins/tensorflow/tf_profiler.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) 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 2d9536c947..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,31 +17,33 @@ """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 from google.cloud.aiplatform.training_utils import environment_variables -from google.cloud.aiplatform.training_utils.cloud_profiler import cloud_profiler_utils from google.cloud.aiplatform.training_utils.cloud_profiler import wsgi_types from google.cloud.aiplatform.training_utils.cloud_profiler.plugins import base_plugin from google.cloud.aiplatform.training_utils.cloud_profiler.plugins.tensorflow import ( tensorboard_api, ) -try: - from tensorboard_plugin_profile.profile_plugin import ProfilePlugin - from werkzeug import Response - import tensorflow as tf -except ImportError as err: - raise ImportError(cloud_profiler_utils.import_error_msg) from err - # TF verison information. Version = namedtuple("Version", ["major", "minor", "patch"])