From e22cfdab34d900091d147f29df4bf8417011de5f Mon Sep 17 00:00:00 2001 From: Will Cromar Date: Thu, 18 Jan 2024 21:06:07 +0000 Subject: [PATCH] use `build_util` for CPU too --- build_util.py | 22 +++++++++++++++ plugins/cpu/README.md | 9 ++---- plugins/cpu/build_util.py | 1 + plugins/cpu/pyproject.toml | 2 +- plugins/cpu/setup.py | 6 ++++ plugins/cpu/torch_xla_cpu_plugin/__init__.py | 2 +- plugins/cuda/README.md | 4 +-- plugins/cuda/setup.py | 29 ++------------------ 8 files changed, 37 insertions(+), 38 deletions(-) create mode 120000 plugins/cpu/build_util.py create mode 100644 plugins/cpu/setup.py diff --git a/build_util.py b/build_util.py index d951d7e31cde..236c86c151b4 100644 --- a/build_util.py +++ b/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)) diff --git a/plugins/cpu/README.md b/plugins/cpu/README.md index b3fedc24d8de..21398769b045 100644 --- a/plugins/cpu/README.md +++ b/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 diff --git a/plugins/cpu/build_util.py b/plugins/cpu/build_util.py new file mode 120000 index 000000000000..219f486130c1 --- /dev/null +++ b/plugins/cpu/build_util.py @@ -0,0 +1 @@ +../../build_util.py \ No newline at end of file diff --git a/plugins/cpu/pyproject.toml b/plugins/cpu/pyproject.toml index 9359a0dc7b5e..770d5cb5fb1e 100644 --- a/plugins/cpu/pyproject.toml +++ b/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" diff --git a/plugins/cpu/setup.py b/plugins/cpu/setup.py new file mode 100644 index 000000000000..9bf28a0ac3b1 --- /dev/null +++ b/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() diff --git a/plugins/cpu/torch_xla_cpu_plugin/__init__.py b/plugins/cpu/torch_xla_cpu_plugin/__init__.py index da7a3234267c..e22f4c65a06c 100644 --- a/plugins/cpu/torch_xla_cpu_plugin/__init__.py +++ b/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 diff --git a/plugins/cuda/README.md b/plugins/cuda/README.md index 2255d09c6188..a8da474948eb 100644 --- a/plugins/cuda/README.md +++ b/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 diff --git a/plugins/cuda/setup.py b/plugins/cuda/setup.py index 486d7a6f21e0..0b6928db9173 100644 --- a/plugins/cuda/setup.py +++ b/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()