Reverts c1e4e41

PiperOrigin-RevId: 690831984

tensorflow/core/framework/BUILD

@@ -844,7 +844,6 @@ tf_cuda_library(
         "//tensorflow/core/runtime_fallback:__subpackages__",
         "//tensorflow/core/tfrt:__subpackages__",
         "//tensorflow/core/util:__pkg__",
-        "//tensorflow/python/profiler/internal:__pkg__",
         "//tensorflow/security/fuzzing:__subpackages__",
     ],
     deps = [
@@ -1104,7 +1103,6 @@ cc_library(
     name = "op",
     srcs = ["op.cc"],
     hdrs = ["op.h"],
-    visibility = default_visibility + ["//tensorflow/python/profiler/internal:__pkg__"],
     deps = [
         ":full_type_proto_cc",
         ":full_type_util",
@@ -1697,7 +1695,6 @@ tf_proto_library(
     visibility = [
         "//tensorflow/core:__subpackages__",
         "//tensorflow/python:__pkg__",
-        "//tensorflow/python/profiler/internal:__pkg__",
         "//tensorflow/security/fuzzing:__subpackages__",
         "//waymo/ml/deploy/benchmark:__subpackages__",
     ],

tensorflow/python/data/experimental/service/server_lib_test.py

@@ -19,6 +19,7 @@
 import threading
 import unittest
 from tensorflow.python.data.experimental.service import server_lib
+from tensorflow.python.framework import errors
 from tensorflow.python.platform import test
 from tensorflow.python.profiler import profiler_client
 
@@ -161,16 +162,13 @@ def testDispatcherNumWorkers(self):
   def testProfileWorker(self):
     dispatcher = server_lib.DispatchServer()
     worker = server_lib.WorkerServer(
-        server_lib.WorkerConfig(dispatcher._address)
-    )
+        server_lib.WorkerConfig(dispatcher._address))
     # Test the profilers are successfully started and connected to profiler
     # service on the worker. Since there is no op running, it is expected to
-    # return RuntimeError with no trace events collected string.
-    with self.assertRaises(RuntimeError) as error:
+    # return UnavailableError with no trace events collected string.
+    with self.assertRaises(errors.UnavailableError) as error:
       profiler_client.trace(worker._address, tempfile.mkdtemp(), duration_ms=10)
-    self.assertStartsWith(
-        str(error.exception), "UNAVAILABLE: No trace event was collected"
-    )
+    self.assertStartsWith(str(error.exception), "No trace event was collected")
 
 
 if __name__ == "__main__":

tensorflow/python/eager/BUILD

@@ -429,7 +429,7 @@ py_strict_library(
         "//tensorflow/core/profiler:internal",
     ],
     deps = [
-        "//tensorflow/python/profiler/internal:_pywrap_profiler_plugin",
+        "//tensorflow/python/profiler/internal:_pywrap_profiler",
         "//tensorflow/python/util:deprecation",
     ],
 )

tensorflow/python/eager/profiler_client.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Profiler client APIs."""
 
-from tensorflow.python.profiler.internal import _pywrap_profiler_plugin
+from tensorflow.python.profiler.internal import _pywrap_profiler
 from tensorflow.python.util.deprecation import deprecated
 
 
@@ -42,21 +42,15 @@ def start_tracing(service_addr,
   Raises:
     UnavailableError: If no trace event is collected.
   """
-  _pywrap_profiler_plugin.trace(
-      service_addr,
-      logdir,
-      worker_list,
-      include_dataset_ops,
-      duration_ms,
-      num_tracing_attempts,
-      {},
-  )
+  _pywrap_profiler.trace(service_addr, logdir, worker_list, include_dataset_ops,
+                         duration_ms, num_tracing_attempts, {})
 
 
 @deprecated('2020-07-01', 'use `tf.profiler.experimental.client.monitor`.')
-def monitor(
-    service_addr, duration_ms, monitoring_level=1, display_timestamp=False
-):
+def monitor(service_addr,
+            duration_ms,
+            monitoring_level=1,
+            display_timestamp=False):
   """Sends grpc requests to profiler server to perform on-demand monitoring.
 
   This method will block caller thread until receives monitoring result.
@@ -71,6 +65,5 @@ def monitor(
   Returns:
     A string of monitoring output.
   """
-  return _pywrap_profiler_plugin.monitor(
-      service_addr, duration_ms, monitoring_level, display_timestamp
-  )
+  return _pywrap_profiler.monitor(service_addr, duration_ms, monitoring_level,
+                                  display_timestamp)

tensorflow/python/eager/profiler_client_test.py

@@ -16,17 +16,18 @@
 
 from tensorflow.python.eager import profiler_client
 from tensorflow.python.eager import test
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 
 
 class ProfilerClientTest(test_util.TensorFlowTestCase):
 
   def testStartTracing_ProcessInvalidAddress(self):
-    with self.assertRaises(RuntimeError):
+    with self.assertRaises(errors.UnavailableError):
       profiler_client.start_tracing('localhost:6006', '/tmp/', 2000)
 
   def testMonitor_ProcessInvalidAddress(self):
-    with self.assertRaises(RuntimeError):
+    with self.assertRaises(errors.UnavailableError):
       profiler_client.monitor('localhost:6006', 2000)
 
 

tensorflow/python/profiler/BUILD

@@ -27,7 +27,7 @@ py_strict_library(
     srcs_version = "PY3",
     deps = [
         "//tensorflow/python/framework:errors",
-        "//tensorflow/python/profiler/internal:_pywrap_profiler_plugin",
+        "//tensorflow/python/profiler/internal:_pywrap_profiler",
         "//tensorflow/python/util:tf_export",
     ],
 )
@@ -88,7 +88,7 @@ py_strict_test(
     deps = [
         "//tensorflow/python/eager:test",
         "//tensorflow/python/framework:test_lib",
-        "//tensorflow/python/profiler/internal:_pywrap_profiler_plugin",
+        "//tensorflow/python/profiler/internal:_pywrap_profiler",
     ],
 )
 

tensorflow/python/profiler/internal/BUILD

@@ -1,5 +1,3 @@
-load("@local_xla//xla/tsl:tsl.bzl", "internal_visibility")
-load("@local_xla//xla/tsl:tsl.default.bzl", "tsl_pybind_extension")
 load("//tensorflow:strict.default.bzl", "py_strict_library", "py_strict_test")
 load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "get_compatible_with_portable", "tf_python_pybind_extension")
 load("//tensorflow/core/profiler/builds:build_config.bzl", "tf_profiler_copts")
@@ -188,58 +186,3 @@ cc_library(
         "@local_xla//xla/tsl/profiler/utils:session_manager",
     ],
 )
-
-tsl_pybind_extension(
-    name = "_pywrap_profiler_plugin",
-    srcs = ["pywrap_profiler_plugin.cc"],
-    pytype_srcs = [
-        "_pywrap_profiler_plugin.pyi",
-    ],
-    visibility = internal_visibility([
-        "//tensorflow/core/profiler:internal",
-        "//tensorflow/python/eager:__pkg__",
-        "//tensorflow/python/profiler:__pkg__",
-    ]),
-    deps = [
-        "//tensorflow/core:protos_all_cc_impl",
-        "//tensorflow/core/framework:attr_value_proto_cc_impl",
-        "//tensorflow/core/framework:op",
-        "//tensorflow/core/framework:tensor",
-        "//tensorflow/core/profiler/convert:repository",
-        "//tensorflow/core/profiler/convert:tool_options",
-        "//tensorflow/core/profiler/convert:xplane_to_tools_data",
-        "//tensorflow/python/lib/core:py_exception_registry",
-        "//tensorflow/python/lib/core:pybind11_status",
-        "@com_google_protobuf//:protobuf",
-        "@local_tsl//tsl/platform:env_impl",
-        "@local_tsl//tsl/profiler/protobuf:profiler_analysis_proto_cc_impl",
-        "@local_tsl//tsl/profiler/protobuf:profiler_options_proto_cc_impl",
-        "@local_tsl//tsl/profiler/protobuf:profiler_service_monitor_result_proto_cc_impl",
-        "@local_tsl//tsl/profiler/protobuf:profiler_service_proto_cc_impl",
-        "@local_xla//xla:autotune_results_proto_cc_impl",
-        "@local_xla//xla:autotuning_proto_cc_impl",
-        "@local_xla//xla:xla_data_proto_cc_impl",
-        "@local_xla//xla:xla_proto_cc_impl",
-        "@local_xla//xla/pjrt:status_casters",
-        "@local_xla//xla/service:hlo_proto_cc_impl",
-        "@local_xla//xla/service/gpu:backend_configs_cc_impl",
-        "@local_xla//xla/service/gpu/model:hlo_op_profile_proto_cc_impl",
-        "@local_xla//xla/stream_executor:device_description_proto_cc_impl",
-        "@local_xla//xla/tsl/framework:allocator_registry_impl",
-        "@local_xla//xla/tsl/lib/io:table",
-        "@local_xla//xla/tsl/profiler/backends/cpu:traceme_recorder_impl",
-        "@local_xla//xla/tsl/profiler/rpc:profiler_server_impl",
-        "@local_xla//xla/tsl/profiler/rpc:profiler_service_impl",
-        "@local_xla//xla/tsl/profiler/rpc/client:capture_profile",
-        "@local_xla//xla/tsl/profiler/rpc/client:profiler_client_impl",
-        "@local_xla//xla/tsl/profiler/utils:session_manager",
-        "@local_xla//xla/tsl/profiler/utils:time_utils_impl",
-        "@local_xla//xla/tsl/protobuf:bfc_memory_map_proto_cc_impl",
-        "@local_xla//xla/tsl/protobuf:dnn_proto_cc_impl",
-        "@local_xla//xla/tsl/protobuf:error_codes_proto_impl_cc_impl",
-        "@local_xla//xla/tsl/protobuf:histogram_proto_cc_impl",
-        "@local_xla//xla/tsl/protobuf:rpc_options_proto_cc_impl",
-        "@local_xla//xla/tsl/protobuf:test_log_proto_cc_impl",
-        "@pybind11",
-    ],
-)

tensorflow/python/profiler/internal/_pywrap_profiler.pyi

@@ -19,4 +19,8 @@ class ProfilerSession:
     def start(self, arg0: str, arg1: dict) -> None: ...
     def stop(self) -> bytes: ...
 
+def monitor(arg0: str, arg1: int, arg2: int, arg3: bool) -> str: ...
 def start_server(arg0: int) -> None: ...
+def trace(arg0: str, arg1: str, arg2: str, arg3: bool, arg4: int, arg5: int, arg6: dict) -> None: ...
+def xspace_to_tools_data(arg0: list, arg1: str, arg2: dict = ...) -> tuple: ...
+def xspace_to_tools_data_from_byte_string(arg0: list, arg1: list, arg2: str, arg3: dict) -> tuple: ...

tensorflow/python/profiler/internal/profiler_pywrap_impl.cc

@@ -47,6 +47,28 @@ namespace pywrap {
 using tsl::profiler::GetRemoteSessionManagerOptionsLocked;
 using tsl::profiler::ValidateHostPortPair;
 
+absl::Status Trace(
+    const char* service_addr, const char* logdir, const char* worker_list,
+    bool include_dataset_ops, int duration_ms, int num_tracing_attempts,
+    const absl::flat_hash_map<std::string, std::variant<int, std::string>>&
+        options) {
+  return tsl::profiler::CaptureRemoteTrace(service_addr, logdir, worker_list,
+                                           include_dataset_ops, duration_ms,
+                                           num_tracing_attempts, options);
+}
+
+absl::Status Monitor(const char* service_addr, int duration_ms,
+                     int monitoring_level, bool display_timestamp,
+                     tensorflow::string* result) {
+  TF_RETURN_IF_ERROR(ValidateHostPortPair(service_addr));
+  {
+    TF_RETURN_IF_ERROR(tsl::profiler::Monitor(service_addr, duration_ms,
+                                              monitoring_level,
+                                              display_timestamp, result));
+  }
+  return absl::OkStatus();
+}
+
 absl::Status ProfilerSessionWrapper::Start(
     const char* logdir,
     const absl::flat_hash_map<std::string, std::variant<int, std::string>>&

tensorflow/python/profiler/internal/profiler_pywrap_impl.h

@@ -27,6 +27,16 @@ namespace tensorflow {
 namespace profiler {
 namespace pywrap {
 
+absl::Status Trace(
+    const char* service_addr, const char* logdir, const char* worker_list,
+    bool include_dataset_ops, int duration_ms, int num_tracing_attempts,
+    const absl::flat_hash_map<std::string, std::variant<int, std::string>>&
+        options);
+
+absl::Status Monitor(const char* service_addr, int duration_ms,
+                     int monitoring_level, bool display_timestamp,
+                     tensorflow::string* result);
+
 class ProfilerSessionWrapper {
  public:
   absl::Status Start(