Add named arguments to the pywrap_profiler_plugin functions.

PiperOrigin-RevId: 851212029

Author: subhamsoni-google

WORKSPACE

@@ -67,10 +67,10 @@ http_archive(
     name = "xla",
     patch_args = ["-p1"],
     patches = ["//third_party:xla.patch"],
-    sha256 = "fb110de1feb871be75ea4ea4c529165ac5feaac61f9a96855d7420f843a61e5d",
-    strip_prefix = "xla-9e7a6358af8fec8d73174bfc6cff2e180bb8fdf6",
+    sha256 = "40c0ab46b7b9f1ccaa063aa5cd3c062e0a7b74e8885250bd6803efc1aa09890b",
+    strip_prefix = "xla-990044eaf27699fc33c874f181c6ff22441b0195",
     urls = [
-        "https://github.com/openxla/xla/archive/9e7a6358af8fec8d73174bfc6cff2e180bb8fdf6.zip",
+        "https://github.com/openxla/xla/archive/990044eaf27699fc33c874f181c6ff22441b0195.zip",
     ],
 )
 

third_party/xla.patch

@@ -1,26 +1,23 @@
 diff --git a/workspace1.bzl b/workspace1.bzl
-index 05121708fc..82e2b0e14d 100644
 --- a/workspace1.bzl
 +++ b/workspace1.bzl
-@@ -2,7 +2,7 @@
-
- load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+@@ -3,7 +3,7 @@
+ 
  load("@com_github_grpc_grpc//bazel:grpc_deps.bzl", "grpc_deps")
 -load("@io_bazel_rules_closure//closure:defs.bzl", "closure_repositories")
 +load("@io_bazel_rules_closure//closure:repositories.bzl", "rules_closure_dependencies")
  load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
+ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
  load("//third_party/llvm:setup.bzl", "llvm_setup")
-
-@@ -13,7 +13,7 @@ def workspace():
+@@ -14,7 +14,7 @@ def workspace():
      native.register_toolchains("@local_config_python//:py_toolchain")
      rules_pkg_dependencies()
-
+ 
 -    closure_repositories()
 +    rules_closure_dependencies()
-
-     http_archive(
+ 
+     tf_http_archive(
          name = "bazel_toolchains",
-
 diff --git a/tools/toolchains/python/python_repo.bzl b/tools/toolchains/python/python_repo.bzl
 index 47fe64d7b7b..a01a1f19c8b 100644
 --- a/tools/toolchains/python/python_repo.bzl

xprof/pywrap/_pywrap_profiler_plugin.pyi

@@ -13,10 +13,29 @@
 # limitations under the License.
 # ==============================================================================
 
-def monitor(arg0: str, arg1: int, arg2: int, arg3: bool) -> str: ...
-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: ...
+def monitor(
+    service_addr: str,
+    duration_ms: int,
+    monitoring_level: int,
+    display_timestamp: bool,
+) -> str: ...
+def trace(
+    service_addr: str,
+    logdir: str,
+    worker_list: str,
+    include_dataset_ops: bool,
+    duration_ms: int,
+    num_tracing_attempts: int,
+    options: dict,
+) -> None: ...
+def xspace_to_tools_data(
+    xspace_path_list: list, tool_name: str, options: dict = ...
+) -> tuple: ...
+def xspace_to_tools_data_from_byte_string(
+    xspace_string_list: list,
+    filenames_list: list,
+    tool_name: str,
+    options: dict,
+) -> tuple: ...
 def start_grpc_server(port: int) -> None: ...
 def initialize_stubs(worker_service_addresses: str) -> None: ...
-

xprof/pywrap/pywrap_profiler_plugin.cc

@@ -57,9 +57,10 @@ ToolOptions ToolOptionsFromPythonDict(const py::dict& dictionary) {
 
 PYBIND11_MODULE(_pywrap_profiler_plugin, m) {
   m.def(
-      "trace", [](const char* service_addr, const char* logdir,
-                  const char* worker_list, bool include_dataset_ops,
-                  int duration_ms, int num_tracing_attempts, py::dict options) {
+      "trace",
+      [](const char* service_addr, const char* logdir, const char* worker_list,
+         bool include_dataset_ops, int duration_ms, int num_tracing_attempts,
+         py::dict options) {
         absl::Status status;
         ToolOptions tool_options = ToolOptionsFromPythonDict(options);
         {
@@ -70,22 +71,29 @@ PYBIND11_MODULE(_pywrap_profiler_plugin, m) {
         }
         // Py_INCREF and Py_DECREF must be called holding the GIL.
         xla::ThrowIfError(status);
-      });
-
-  m.def("monitor", [](const char* service_addr, int duration_ms,
-                      int monitoring_level, bool display_timestamp) {
-    std::string content;
-    absl::Status status;
-    {
-      py::gil_scoped_release release;
-      status =
-          xprof::pywrap::Monitor(service_addr, duration_ms, monitoring_level,
-                                 display_timestamp, &content);
-    }
-    // Py_INCREF and Py_DECREF must be called holding the GIL.
-    xla::ThrowIfError(status);
-    return content;
-  });
+      },
+      py::arg("service_addr"), py::arg("logdir"), py::arg("worker_list"),
+      py::arg("include_dataset_ops"), py::arg("duration_ms"),
+      py::arg("num_tracing_attempts"), py::arg("options"));
+
+  m.def(
+      "monitor",
+      [](const char* service_addr, int duration_ms, int monitoring_level,
+         bool display_timestamp) {
+        std::string content;
+        absl::Status status;
+        {
+          py::gil_scoped_release release;
+          status = xprof::pywrap::Monitor(service_addr, duration_ms,
+                                          monitoring_level, display_timestamp,
+                                          &content);
+        }
+        // Py_INCREF and Py_DECREF must be called holding the GIL.
+        xla::ThrowIfError(status);
+        return content;
+      },
+      py::arg("service_addr"), py::arg("duration_ms"),
+      py::arg("monitoring_level"), py::arg("display_timestamp"));
 
   m.def(
       "xspace_to_tools_data",
@@ -98,8 +106,7 @@ PYBIND11_MODULE(_pywrap_profiler_plugin, m) {
           xspace_paths.push_back(xspace_path);
         }
         std::string tool_name = std::string(py_tool_name);
-        ToolOptions tool_options =
-            ToolOptionsFromPythonDict(options);
+        ToolOptions tool_options = ToolOptionsFromPythonDict(options);
         absl::StatusOr<std::pair<std::string, bool>> result;
         {
           py::gil_scoped_release release;
@@ -113,7 +120,8 @@ PYBIND11_MODULE(_pywrap_profiler_plugin, m) {
         return py::make_tuple(py::bytes(result->first),
                               py::bool_(result->second));
       },
-      py::arg(), py::arg(), py::arg() = py::dict());
+      py::arg("xspace_path_list"), py::arg("tool_name"),
+      py::arg("options") = py::dict());
 
   m.def(
       "xspace_to_tools_data_from_byte_string",
@@ -132,8 +140,7 @@ PYBIND11_MODULE(_pywrap_profiler_plugin, m) {
         }
 
         std::string tool_name = std::string(py_tool_name);
-        ToolOptions tool_options =
-            ToolOptionsFromPythonDict(options);
+        ToolOptions tool_options = ToolOptionsFromPythonDict(options);
 
         absl::StatusOr<std::pair<std::string, bool>> result;
         {
@@ -148,17 +155,24 @@ PYBIND11_MODULE(_pywrap_profiler_plugin, m) {
         return py::make_tuple(py::bytes(result->first),
                               py::bool_(result->second));
       },
-      py::arg(), py::arg(), py::arg(), py::arg() = py::dict());
+      py::arg("xspace_string_list"), py::arg("filenames_list"),
+      py::arg("tool_name"), py::arg("options") = py::dict());
 
-  m.def("start_grpc_server", [](int port) {
-    py::gil_scoped_release release;
-    xprof::pywrap::StartGrpcServer(port);
-  });
+  m.def(
+      "start_grpc_server",
+      [](int port) {
+        py::gil_scoped_release release;
+        xprof::pywrap::StartGrpcServer(port);
+      },
+      py::arg("port"));
 
-  m.def("initialize_stubs", [](const std::string& worker_service_addresses) {
-    py::gil_scoped_release release;
-    xprof::profiler::InitializeStubs(worker_service_addresses);
-  });
+  m.def(
+      "initialize_stubs",
+      [](const std::string& worker_service_addresses) {
+        py::gil_scoped_release release;
+        xprof::profiler::InitializeStubs(worker_service_addresses);
+      },
+      py::arg("worker_service_addresses"));
 };
 
 }  // namespace