feat: Add profiling initialization code to training_utils

google/cloud/aiplatform/tensorboard/__init__.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 #
 
-from google.cloud.aiplatform.tensorboard.tensorboard import Tensorboard
+from google.cloud.aiplatform.tensorboard.tensorboard_resource import Tensorboard
 
 
 __all__ = ("Tensorboard",)

google/cloud/aiplatform/tensorboard/uploader_utils.py

@@ -318,7 +318,6 @@ def get_or_create(
                         filter="display_name = {}".format(json.dumps(str(tag_name))),
                     )
                 )
-
                 num = 0
                 time_series = None
 

google/cloud/aiplatform/training_utils/__init__.py

@@ -0,0 +1,22 @@
+# -*- 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.
+#
+
+from google.cloud.aiplatform.training_utils.environment_variables import (
+    EnvironmentVariables,
+)
+
+__all__ = ("EnvironmentVariables",)

google/cloud/aiplatform/training_utils/cloud_profiler/README.rst

@@ -0,0 +1,20 @@
+Cloud Profiler
+=================================
+
+Cloud Profiler allows you to profile your remote Vertex AI Training jobs on demand. To understand the performance of your training code, a user can update their scripts by initializing the profiler and capturing the profile session through Vertex TensorBoard.
+
+Quick Start
+------------
+
+To start using the profiler, update the training script to include the following:
+
+.. code-block:: Python
+
+    from google.cloud.aiplatform.training_utils import cloud_profiler
+    ...
+    cloud_profiler.init()
+
+
+Next, run the job with with a Vertex TensorBoard instance. For full details on how to do this, visit https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-overview
+
+Finally, visit your TensorBoard instance and click the `Capture Profile` button. This will allow users to capture profiling statistics for the running jobs.

google/cloud/aiplatform/training_utils/cloud_profiler/__init__.py

@@ -0,0 +1,29 @@
+# -*- 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.
+#
+
+from google.cloud.aiplatform.training_utils.cloud_profiler import initializer
+
+"""
+Initialize the cloud profiler for tensorflow.
+
+Usage:
+from google.cloud.aiplatform.training_utils import cloud_profiler
+
+cloud_profiler.init(profiler='tensorflow')
+"""
+
+init = initializer.initialize

google/cloud/aiplatform/training_utils/cloud_profiler/base_plugin.py

@@ -0,0 +1,70 @@
+# -*- 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 abc
+from typing import Dict
+
+
+class BasePlugin(abc.ABC):
+    """Base plugin for cloud training tools endpoints.
+
+    The plugins support registering http handlers to be used for
+    AI Platform training jobs.
+    """
+
+    @staticmethod
+    @abc.abstractmethod
+    def setup() -> None:
+        """Run any setup code for the plugin before webserver is launched."""
+        raise NotImplementedError
+
+    @staticmethod
+    @abc.abstractmethod
+    def can_initialize() -> bool:
+        """Check whether a plugin is able to be initialized.
+
+        Used for checking if correct dependencies are installed, system requirements, etc.
+
+        Returns:
+            Bool indicating whether the plugin can be initialized.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abc.abstractmethod
+    def post_setup_check() -> bool:
+        """Check if after initialization, we need to use the plugin.
+
+        Example: Web server only needs to run for main node for training, others
+        just need to have 'setup()' run to start the rpc server.
+
+        Returns:
+            A boolean indicating whether post setup checks pass.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_routes(self) -> Dict[str, str]:
+        """Get the mapping from path to handler.
+
+        This is the method in which plugins can assign different routes to
+        different handlers.
+
+        Returns:
+            A Dict[str, str] mapping a route to a handler.
+        """
+        raise NotImplementedError

google/cloud/aiplatform/training_utils/cloud_profiler/initializer.py

@@ -0,0 +1,100 @@
+# -*- 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 logging
+import threading
+from typing import Callable
+from werkzeug import serving
+
+from google.cloud.aiplatform.training_utils.cloud_profiler import base_plugin
+from google.cloud.aiplatform.training_utils.cloud_profiler import webserver
+from google.cloud.aiplatform.training_utils.cloud_profiler.plugins import tf_profiler
+
+
+_AVAILABLE_PLUGINS = {"tensorflow": tf_profiler.TFProfiler}
+_HOST_PORT = 6010
+
+
+def _build_plugin(
+    plugin: Callable[[], base_plugin.BasePlugin]
+) -> base_plugin.BasePlugin:
+    """Builds the plugin given the object.
+
+    Args:
+        plugin (Callable[[], base_plugin):
+            Required. An uninitialized plugin.
+
+    Returns:
+        An initialized plugin.
+    """
+    if not plugin.can_initialize():
+        logging.warning("Cannot initialize the plugin")
+        return
+
+    plugin.setup()
+
+    if not plugin.post_setup_check():
+        return
+
+    return plugin()
+
+
+def _run_app_thread(server: webserver.WebServer):
+    """Run the webserver in a separate thread.
+
+    Args:
+        server (webserver.WebServer):
+            Required. A webserver to accept requests.
+    """
+    daemon = threading.Thread(
+        name="profile_server",
+        target=serving.run_simple,
+        args=("0.0.0.0", _HOST_PORT, server,),
+    )
+    daemon.setDaemon(True)
+    daemon.start()
+
+
+def initialize(plugin: str = "tensorflow"):
+    """Initializes the profiling SDK.
+
+    Args:
+        plugin (str):
+            Required. Name of the plugin to initialize.
+            Current options are ["tensorflow"]
+
+    Raises:
+        ValueError:
+            The plugin does not exist.
+    """
+
+    plugin_obj = _AVAILABLE_PLUGINS.get(plugin)
+
+    if not plugin_obj:
+        raise ValueError(
+            "Plugin {} not available, must choose from {}".format(
+                plugin, _AVAILABLE_PLUGINS.keys()
+            )
+        )
+
+    prof_plugin = _build_plugin(plugin_obj)
+
+    if not prof_plugin:
+        return
+
+    server = webserver.WebServer([prof_plugin])
+    _run_app_thread(server)