Fix XProf server not coming up on windows machine

PiperOrigin-RevId: 789880022

Author: Profiler Team

plugin/build_pip_package.sh

@@ -84,6 +84,10 @@ find . -name "*.bak" -exec rm {} \;
 
 cp ${build_workspace}/bazel-bin/xprof/pywrap/_pywrap_profiler_plugin.so xprof/convert/
 
+if [[ "$(uname)" == *MSYS_NT* ]]; then
+  mv xprof/convert/_pywrap_profiler_plugin.so xprof/convert/_pywrap_profiler_plugin.pyd
+fi
+
 # Copy static files.
 cd xprof
 mkdir -p static

plugin/setup.py

@@ -59,7 +59,7 @@ def get_readme():
     ),
     package_data={
         'xprof': ['static/**'],
-        '': ['_pywrap_profiler_plugin.so'],
+        '': ['_pywrap_profiler_plugin.so', '_pywrap_profiler_plugin.pyd'],
     },
     entry_points={
         'tensorboard_plugins': [

plugin/xprof/standalone/BUILD

@@ -2,6 +2,7 @@
 #   A collection of classes to bypass Tensorboard dependencies.
 
 load("@python_deps//:requirements.bzl", "requirement")
+load("@rules_python//python:defs.bzl", "py_test")
 load("@rules_python//python:py_library.bzl", "py_library")
 
 visibility = ["//plugin:internal"]
@@ -66,3 +67,13 @@ py_library(
     srcs_version = "PY2AND3",
     deps = [],
 )
+
+py_test(
+    name = "plugin_asset_util_test",
+    srcs = ["plugin_asset_util_test.py"],
+    deps = [
+        ":plugin_asset_util",
+        "//testing/pybase",
+        "//testing/pybase:parameterized",
+    ],
+)

plugin/xprof/standalone/plugin_asset_util.py

@@ -23,10 +23,29 @@
 _PROTOCOL_PREFIXES = ("gs://",)
 
 
+def _get_protocol_from_parsed_url_scheme(scheme: str) -> str:
+  """Returns the protocol from the parsed URL scheme.
+
+  If the scheme is not recognized, returns "file" as the default protocol.
+
+  Args:
+    scheme: The parsed URL scheme.
+
+  Returns:
+    The protocol from the parsed URL scheme.
+  """
+  if not scheme:
+    return "file"
+  for prefix in _PROTOCOL_PREFIXES:
+    if prefix.startswith(scheme):
+      return scheme
+  return "file"
+
+
 def get_fs_protocol(path):
   string_path = str(path)
   parsed_url = urllib.parse.urlparse(string_path)
-  protocol = parsed_url.scheme or "file"
+  protocol = _get_protocol_from_parsed_url_scheme(parsed_url.scheme)
   return fsspec.filesystem(protocol)
 
 

plugin/xprof/standalone/plugin_asset_util_test.py

@@ -0,0 +1,36 @@
+from google3.testing.pybase import googletest
+from google3.testing.pybase import parameterized
+from xprof.standalone import plugin_asset_util
+
+
+class PluginAssetUtilTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="No_Scheme",
+          scheme="",
+          expected="file",
+      ),
+      dict(
+          testcase_name="GCS_Scheme",
+          scheme="gs",
+          expected="gs",
+      ),
+      dict(
+          testcase_name="Windows_Path_C_Drive",
+          scheme="C",
+          expected="file",
+      ),
+      dict(
+          testcase_name="Windows_path_T_Drive",
+          scheme="T",
+          expected="file",
+      ),
+  )
+  def test_get_protocol_from_parsed_url_scheme(self, scheme, expected):
+    got = plugin_asset_util._get_protocol_from_parsed_url_scheme(scheme=scheme)
+    self.assertEqual(expected, got)
+
+
+if __name__ == "__main__":
+  googletest.main()