Add subprocess registration hooks to JAX profiler libraries.

PiperOrigin-RevId: 893734302

Author: bmass02

xla/python/BUILD

@@ -458,6 +458,7 @@ nanobind_pywrap_extension(
         "//xla/backends/profiler:profiler_backends",
         "//xla/backends/profiler/cpu:python_tracer",
         "//xla/backends/profiler/plugin:plugin_tracer",
+        "//xla/backends/profiler/subprocess:subprocess_registry",
         "//xla/backends/profiler/util:metadata_registry",
         "//xla/pjrt:exceptions",
         "//xla/pjrt:status_casters",

xla/python/_profiler.pyi

@@ -24,6 +24,8 @@ class ProfilerServer:
 
 def start_server(port: int) -> ProfilerServer: ...
 def register_plugin_profiler(arg: CapsuleType, /) -> None: ...
+def register_subprocess(pid: int, port: int) -> bool: ...
+def unregister_subprocess(pid: int) -> bool: ...
 
 class ProfilerSession:
   @overload

xla/python/profiler.cc

@@ -28,6 +28,7 @@ limitations under the License.
 #include "nanobind/stl/string_view.h"  // IWYU pragma: keep
 #include "nanobind/stl/unique_ptr.h"  // IWYU pragma: keep
 #include "nanobind/stl/vector.h"  // IWYU pragma: keep
+#include "xla/backends/profiler/subprocess/subprocess_registry.h"
 #include "xla/backends/profiler/util/metadata_registry.h"
 #include "xla/pjrt/c/pjrt_c_api.h"
 #include "xla/pjrt/exceptions.h"
@@ -176,6 +177,18 @@ NB_MODULE(_profiler, m) {
         ") -> None"
           // clang-format on
           ));
+  m.def(
+      "register_subprocess",
+      [](int pid, int port) -> bool {
+        return xla::profiler::subprocess::RegisterSubprocess(pid, port).ok();
+      },
+      nb::arg("pid"), nb::arg("port"));
+  m.def(
+      "unregister_subprocess",
+      [](int pid) -> bool {
+        return xla::profiler::subprocess::UnregisterSubprocess(pid).ok();
+      },
+      nb::arg("pid"));
 
   nb::class_<ProfilerSessionWrapper> profiler_session_class(m,
                                                             "ProfilerSession");