Export public symbols for programmatic profiling APIs.

PiperOrigin-RevId: 296053405
Change-Id: I2d52a45e61bcb37d39d45be89c26a6a3f8f3ff1b

tensorflow/python/BUILD

@@ -204,6 +204,7 @@ py_library(
         "//tensorflow/python/ops/ragged",
         "//tensorflow/python/ops/signal",
         "//tensorflow/python/profiler",
+        "//tensorflow/python/profiler:profiler_v2",
         "//tensorflow/python/saved_model",
         "//tensorflow/python/tools:module_util",
         "//tensorflow/python/tools/api/generator:create_python_api",

tensorflow/python/__init__.py

@@ -111,6 +111,7 @@ from tensorflow.python.ops.linalg.sparse import sparse
 from tensorflow.python.ops.losses import losses
 from tensorflow.python.ops.signal import signal
 from tensorflow.python.profiler import profiler
+from tensorflow.python.profiler import profiler_v2
 from tensorflow.python.saved_model import saved_model
 from tensorflow.python.summary import summary
 from tensorflow.python.tpu import api

tensorflow/python/profiler/profiler_v2.py

@@ -39,11 +39,13 @@ import threading
 from tensorflow.python.framework import errors
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.profiler.internal import _pywrap_profiler
+from tensorflow.python.util.tf_export import tf_export
 
 _profiler = None
 _profiler_lock = threading.Lock()
 
 
+@tf_export('profiler.experimental.start', v1=[])
 def start(logdir):
   """Starts profiling.
 
@@ -55,9 +57,9 @@ def start(logdir):
 
   Example usage:
   ```python
-  tf.profiler.start('logdir_path')
+  tf.profiler.experimental.start('logdir_path')
   # do your training here.
-  tf.profiler.stop()
+  tf.profiler.experimental.stop()
   ```
 
   Launch TensorBoard and point it to the same logdir you provided to this API.
@@ -81,10 +83,11 @@ def start(logdir):
                                       'Another profiler is running.')
 
 
+@tf_export('profiler.experimental.stop', v1=[])
 def stop(save=True):
   """Stops the current profiling session.
 
-  The profiler session will be stopped and profile results will be saved.
+  The profiler session will be stopped and profile results can be saved.
 
   Args:
     save: An optional variable to save the results to TensorBoard. Default True.
@@ -103,6 +106,7 @@ def stop(save=True):
     _profiler = None
 
 
+@tf_export('profiler.experimental.server.start', v1=[])
 def start_server(port):
   """Start a profiler grpc server that listens to given port.
 
@@ -111,16 +115,26 @@ def start_server(port):
 
   Args:
     port: port profiler server listens to.
+
+  Example usage:
+  ```python
+  tf.profiler.experimental.server.start('6009')
+  # do your training here.
+
   """
   _pywrap_profiler.start_server(port)
 
 
-class Profiler(object):
-  """Context-manager profiler API.
+@tf_export('profiler.experimental.Profile', v1=[])
+class Profile(object):
+  """Context-manager profile API.
+
+  Profiling will start when entering the scope, and stop and save the results to
+  the logdir when exits the scope. Open TensorBoard profile tab to view results.
 
   Example usage:
   ```python
-  with Profiler("/path/to/logdir"):
+  with tf.profiler.experimental.Profile("/path/to/logdir"):
     # do some work
   ```
   """

tensorflow/python/tools/api/generator/api_init_files.bzl

@@ -48,6 +48,9 @@ TENSORFLOW_API_INIT_FILES = [
     "mlir/experimental/__init__.py",
     "nest/__init__.py",
     "nn/__init__.py",
+    "profiler/__init__.py",
+    "profiler/experimental/__init__.py",
+    "profiler/experimental/server/__init__.py",
     "quantization/__init__.py",
     "ragged/__init__.py",
     "random/__init__.py",

tensorflow/tools/api/golden/v2/tensorflow.pbtxt

@@ -292,6 +292,10 @@ tf_module {
     name: "optimizers"
     mtype: "<type \'module\'>"
   }
+  member {
+    name: "profiler"
+    mtype: "<type \'module\'>"
+  }
   member {
     name: "qint16"
     mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"

tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profile.pbtxt

@@ -0,0 +1,9 @@
+path: "tensorflow.profiler.experimental.Profile"
+tf_class {
+  is_instance: "<class \'tensorflow.python.profiler.profiler_v2.Profile\'>"
+  is_instance: "<type \'object\'>"
+  member_method {
+    name: "__init__"
+    argspec: "args=[\'self\', \'logdir\'], varargs=None, keywords=None, defaults=None"
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.pbtxt

@@ -0,0 +1,19 @@
+path: "tensorflow.profiler.experimental"
+tf_module {
+  member {
+    name: "Profile"
+    mtype: "<type \'type\'>"
+  }
+  member {
+    name: "server"
+    mtype: "<type \'module\'>"
+  }
+  member_method {
+    name: "start"
+    argspec: "args=[\'logdir\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "stop"
+    argspec: "args=[\'save\'], varargs=None, keywords=None, defaults=[\'True\'], "
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.server.pbtxt

@@ -0,0 +1,7 @@
+path: "tensorflow.profiler.experimental.server"
+tf_module {
+  member_method {
+    name: "start"
+    argspec: "args=[\'port\'], varargs=None, keywords=None, defaults=None"
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.profiler.pbtxt

@@ -0,0 +1,7 @@
+path: "tensorflow.profiler"
+tf_module {
+  member {
+    name: "experimental"
+    mtype: "<type \'module\'>"
+  }
+}