[bazel] add python targets (#101003)
This PR adds bazel python, so that bazel build could be used from python like `import torch`.
Notable changes:
- Add the python targets.
- Add the version.py.tpl generation.
- In order to archive the `USE_GLOBAL_DEPS = False` just for the bazel build, employ a monkey-patch hack in the mentioned `version.py.tpl`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/101003
Approved by: https://github.com/huydhn
diff --git a/.bazelrc b/.bazelrc
index f0b41ec..7581b52 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -96,6 +96,9 @@
# likely want to have this disabled for the most part.
build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-missing-field-initializers
+build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-unused-function
+build --per_file_copt='^//.*\.(cpp|cc)$'@-Wno-unused-variable
+
build --per_file_copt='//:aten/src/ATen/RegisterCompositeExplicitAutograd\.cpp$'@-Wno-error=unused-function
build --per_file_copt='//:aten/src/ATen/RegisterCompositeImplicitAutograd\.cpp$'@-Wno-error=unused-function
build --per_file_copt='//:aten/src/ATen/RegisterMkldnnCPU\.cpp$'@-Wno-error=unused-function
diff --git a/.ci/pytorch/build.sh b/.ci/pytorch/build.sh
index da173b1..ead26a7 100755
--- a/.ci/pytorch/build.sh
+++ b/.ci/pytorch/build.sh
@@ -197,7 +197,7 @@
if [[ "$CUDA_VERSION" == "cpu" ]]; then
# Build torch, the Python module, and tests for CPU-only
- tools/bazel build --config=no-tty "${BAZEL_MEM_LIMIT}" "${BAZEL_CPU_LIMIT}" --config=cpu-only :torch :_C.so :all_tests
+ tools/bazel build --config=no-tty "${BAZEL_MEM_LIMIT}" "${BAZEL_CPU_LIMIT}" --config=cpu-only :torch :torch/_C.so :all_tests
else
tools/bazel build --config=no-tty "${BAZEL_MEM_LIMIT}" "${BAZEL_CPU_LIMIT}" //...
fi
diff --git a/.ci/pytorch/test.sh b/.ci/pytorch/test.sh
index 69ba42e..563972e 100755
--- a/.ci/pytorch/test.sh
+++ b/.ci/pytorch/test.sh
@@ -889,6 +889,7 @@
//:torch_dist_autograd_test \
//:torch_include_test \
//:transformer_test \
+ //:test_bazel \
//c10/cuda/test:test \
//c10/test:core_tests \
//c10/test:typeid_test \
diff --git a/BUILD.bazel b/BUILD.bazel
index 84aa029..32a2e19 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -1,14 +1,15 @@
load("@bazel_skylib//lib:paths.bzl", "paths")
load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_test")
-load("@pytorch//third_party:substitution.bzl", "header_template_rule")
+load("@pytorch//third_party:substitution.bzl", "header_template_rule", "template_rule")
load("@pytorch//:tools/bazel.bzl", "rules")
load("@pytorch//tools/rules:cu.bzl", "cu_library")
load("@pytorch//tools/config:defs.bzl", "if_cuda")
load("@pytorch//:aten.bzl", "generate_aten", "intern_build_aten_ops")
load(":build.bzl", "GENERATED_AUTOGRAD_CPP", "GENERATED_AUTOGRAD_PYTHON", "define_targets")
-load(":build_variables.bzl", "jit_core_sources", "lazy_tensor_ts_sources", "libtorch_core_sources", "libtorch_cuda_sources", "libtorch_distributed_sources", "libtorch_extra_sources", "libtorch_nvfuser_generated_headers", "libtorch_nvfuser_runtime_sources", "libtorch_python_core_sources", "torch_cpp_srcs")
+load(":build_variables.bzl", "jit_core_sources", "lazy_tensor_ts_sources", "libtorch_core_sources", "libtorch_cuda_sources", "libtorch_distributed_sources", "libtorch_extra_sources", "libtorch_nvfuser_generated_headers", "libtorch_nvfuser_runtime_sources", "libtorch_python_core_sources", "torch_cpp_srcs", "libtorch_python_cuda_sources", "libtorch_python_distributed_sources")
load(":ufunc_defs.bzl", "aten_ufunc_generated_cpu_kernel_sources", "aten_ufunc_generated_cpu_sources", "aten_ufunc_generated_cuda_sources")
+load("//:tools/bazel.bzl", "rules")
define_targets(rules = rules)
@@ -1671,19 +1672,72 @@
cc_library(
name = "torch_python",
- srcs = libtorch_python_core_sources + GENERATED_AUTOGRAD_PYTHON,
+ srcs = libtorch_python_core_sources
+ + if_cuda(libtorch_python_cuda_sources)
+ + if_cuda(libtorch_python_distributed_sources)
+ + GENERATED_AUTOGRAD_PYTHON,
+ hdrs = glob([
+ "torch/csrc/generic/*.cpp",
+ ]),
+ copts = COMMON_COPTS + if_cuda(["-DUSE_CUDA=1"]),
deps = [
- ":shm",
":torch",
+ ":shm",
"@pybind11",
],
)
pybind_extension(
- name = "_C",
+ name = "torch/_C",
srcs = ["torch/csrc/stub.c"],
deps = [
":torch_python",
+ ":aten_nvrtc",
+ ],
+)
+
+cc_binary(
+ name = "torch/bin/torch_shm_manager",
+ srcs = [
+ "torch/lib/libshm/manager.cpp",
+ ],
+ deps = [
+ ":shm",
+ ],
+ linkstatic = False,
+)
+
+template_rule(
+ name = "gen_version_py",
+ src = ":torch/version.py.tpl",
+ out = "torch/version.py",
+ substitutions = if_cuda({
+ # Set default to 11.2. Otherwise Torchvision complains about incompatibility.
+ "{{CUDA_VERSION}}": "11.2",
+ "{{VERSION}}": "2.0.0",
+ }, {
+ "{{CUDA_VERSION}}": "None",
+ "{{VERSION}}": "2.0.0",
+ }),
+)
+
+rules.py_library(
+ name = "pytorch_py",
+ visibility = ["//visibility:public"],
+ srcs = glob(["torch/**/*.py"], exclude = ["torch/version.py"]) + [":torch/version.py"],
+ deps = [
+ rules.requirement("future"),
+ rules.requirement("numpy"),
+ rules.requirement("pyyaml"),
+ rules.requirement("requests"),
+ rules.requirement("setuptools"),
+ rules.requirement("six"),
+ rules.requirement("typing_extensions"),
+ "//torchgen",
+ ],
+ data = [
+ ":torch/_C.so",
+ ":torch/bin/torch_shm_manager",
],
)
@@ -1845,6 +1899,14 @@
],
)
+# python api tests
+
+py_test(
+ name = "test_bazel",
+ srcs = ["test/test_bazel.py"],
+ deps = [":pytorch_py"],
+)
+
# all tests
test_suite(
name = "all_tests",
diff --git a/WORKSPACE b/WORKSPACE
index 9272e44..05d0746 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -54,6 +54,13 @@
urls = [
"https://github.com/google/glog/archive/v0.4.0.tar.gz",
],
+ build_file_content = """
+licenses(['notice'])
+
+load(':bazel/glog.bzl', 'glog_library')
+# TODO: figure out why enabling gflags leads to SIGSEV on the logging init
+glog_library(with_gflags=0)
+ """
)
http_archive(
diff --git a/test/test_bazel.py b/test/test_bazel.py
new file mode 100644
index 0000000..645f7f9
--- /dev/null
+++ b/test/test_bazel.py
@@ -0,0 +1,15 @@
+# Owner(s): ["module: bazel"]
+
+"""
+This test module contains a minimalistic "smoke tests" for the bazel build.
+
+Currently it doesn't use any testing framework (i.e. pytest)
+TODO: integrate this into the existing pytorch testing framework.
+"""
+
+import torch
+
+def test_sum():
+ assert torch.eq(torch.tensor([[1, 2, 3]]) + torch.tensor([[4, 5, 6]]), torch.tensor([[5, 7, 9]])).all()
+
+test_sum()
diff --git a/torch/version.py.tpl b/torch/version.py.tpl
new file mode 100644
index 0000000..26b80ac
--- /dev/null
+++ b/torch/version.py.tpl
@@ -0,0 +1,10 @@
+__version__ = '{{VERSION}}'
+debug = False
+cuda = '{{CUDA_VERSION}}'
+hip = None
+
+# This is a gross monkey-patch hack that depends on the order of imports
+# in torch/__init__.py
+# TODO: find a more elegant solution to set `USE_GLOBAL_DEPS` for the bazel build
+import torch
+torch.USE_GLOBAL_DEPS = False
