use build_util for CPU too

build_util.py

@@ -1,5 +1,8 @@
 import os
 from typing import Iterable
+import subprocess
+import sys
+import shutil
 
 def check_env_flag(name: str, default: str = '') -> bool:
   return os.getenv(name, default).upper() in ['ON', '1', 'YES', 'TRUE', 'Y']
@@ -49,3 +52,22 @@ def bazel_options_from_env() -> Iterable[str]:
     bazel_flags.append('--config=acl')
 
   return bazel_flags
+
+def bazel_build(bazel_target: str, destination_dir: str, options: Iterable[str] = []):
+  bazel_argv = ['bazel', 'build', bazel_target, f"--symlink_prefix={os.path.join(os.getcwd(), 'bazel-')}"]
+
+  # Remove duplicated flags because they confuse bazel
+  flags = set(bazel_options_from_env() + options)
+  bazel_argv.extend(flags)
+
+  print(' '.join(bazel_argv), flush=True)
+  subprocess.check_call(bazel_argv, stdout=sys.stdout, stderr=sys.stderr)
+
+  target_path = bazel_target.replace('@xla//', 'external/xla/').replace('//', '').replace(':', '/')
+  output_path = os.path.join('bazel-bin', target_path)
+  output_filename = os.path.basename(output_path)
+
+  if not os.path.exists(destination_dir):
+      os.makedirs(destination_dir)
+
+  shutil.copyfile(output_path, os.path.join(destination_dir, output_filename))

plugins/cpu/README.md

@@ -10,15 +10,10 @@ repository (see `bazel build` command below).
 ## Building
 
 ```bash
-# Build PJRT plugin
-bazel build //plugins/cpu:pjrt_c_api_cpu_plugin.so --cxxopt=-D_GLIBCXX_USE_CXX11_ABI=1
-# Copy to package dir
-cp bazel-bin/plugins/cpu/pjrt_c_api_cpu_plugin.so plugins/cpu/torch_xla_cpu_plugin/
-
 # Build wheel
-pip wheel plugins/cpu
+pip wheel plugins/cpu --no-build-isolation -v
 # Or install directly
-pip install plugins/cpu
+pip install plugins/cpu --no-build-isolation -v
 ```
 
 ## Usage

plugins/cpu/build_util.py

@@ -0,0 +1 @@
+../../build_util.py

plugins/cpu/pyproject.toml

@@ -12,7 +12,7 @@ description = "CPU PJRT Plugin for testing only"
 requires-python = ">=3.8"
 
 [tool.setuptools.package-data]
-torch_xla_cpu_plugin = ["*.so"]
+torch_xla_cpu_plugin = ["lib/*.so"]
 
 [project.entry-points."torch_xla.plugins"]
 cpu = "torch_xla_cpu_plugin:CpuPlugin"

plugins/cpu/setup.py

@@ -0,0 +1,6 @@
+import build_util
+import setuptools
+
+build_util.bazel_build('//plugins/cpu:pjrt_c_api_cpu_plugin.so', 'torch_xla_cpu_plugin/lib')
+
+setuptools.setup()

plugins/cpu/torch_xla_cpu_plugin/__init__.py

@@ -4,7 +4,7 @@
 
 class CpuPlugin(plugins.DevicePlugin):
   def library_path(self) -> str:
-    return os.path.join(os.path.dirname(__file__), 'pjrt_c_api_cpu_plugin.so')
+    return os.path.join(os.path.dirname(__file__), 'lib/pjrt_c_api_cpu_plugin.so')
 
   def physical_chip_count(self) -> int:
     return 1

plugins/cuda/README.md

@@ -8,9 +8,9 @@ repository (see `bazel build` command below).
 
 ```bash
 # Build wheel
-pip wheel plugins/cuda
+pip wheel plugins/cuda --no-build-isolation -v
 # Or install directly
-pip install plugins/cuda
+pip install plugins/cuda --no-build-isolation -v
 ```
 
 ## Usage

plugins/cuda/setup.py

@@ -1,31 +1,6 @@
-import subprocess
-import sys
-from typing import Iterable
-import setuptools
-import shutil
-import os
-
 import build_util
+import setuptools
 
-def _bazel_build(bazel_target: str, destination_path: str, options: Iterable[str] = []):
-  bazel_argv = ['bazel', 'build', bazel_target, f"--symlink_prefix={os.path.join(os.getcwd(), 'bazel-')}"]
-
-  # Remove duplicated flags because they confuse bazel
-  flags = set(build_util.bazel_options_from_env() + options)
-  bazel_argv.extend(flags)
-
-  print(' '.join(bazel_argv), flush=True)
-  subprocess.check_call(bazel_argv, stdout=sys.stdout, stderr=sys.stderr)
-
-  target_path = bazel_target.replace('@xla//', 'external/xla/').replace(':', '/')
-  output_path = os.path.join('bazel-bin', target_path)
-  output_filename = os.path.basename(output_path)
-
-  if not os.path.exists(destination_path):
-      os.makedirs(destination_path)
-
-  shutil.copyfile(output_path, os.path.join(destination_path, output_filename))
-
-_bazel_build('@xla//xla/pjrt/c:pjrt_c_api_gpu_plugin.so', 'torch_xla_cuda_plugin/lib', ['--config=cuda'])
+build_util.bazel_build('@xla//xla/pjrt/c:pjrt_c_api_gpu_plugin.so', 'torch_xla_cuda_plugin/lib', ['--config=cuda'])
 
 setuptools.setup()