160 lines
3.7 KiB
Python
160 lines
3.7 KiB
Python
load("@bazel_skylib//rules:select_file.bzl", "select_file")
|
|
load("@rules_cc//cc:cc_library.bzl", "cc_library")
|
|
load("@rules_cc//cc:cc_shared_library.bzl", "cc_shared_library")
|
|
load("@rules_python//python:pip.bzl", "compile_pip_requirements")
|
|
load("@rules_python//python/entry_points:py_console_script_binary.bzl", "py_console_script_binary")
|
|
load("@rules_uv//uv:pip.bzl", "pip_compile")
|
|
load("@rules_zig//zig:defs.bzl", "zig_library")
|
|
load("@zml//bazel:cc_import.bzl", "cc_import")
|
|
load("@zml//bazel:runfiles.bzl", "runfiles_to_default")
|
|
load(":neuron.bzl", "py_binary_with_script")
|
|
load(":pyenv.bzl", "pyenv_zig")
|
|
|
|
pyenv_zig(
|
|
name = "libneuronxla_pyenv_zig",
|
|
deps = [":libneuronxla"],
|
|
)
|
|
|
|
zig_library(
|
|
name = "libneuronxla_pyenv",
|
|
main = ":libneuronxla_pyenv_zig",
|
|
)
|
|
|
|
pip_compile(
|
|
name = "update_requirements",
|
|
args = [
|
|
"--generate-hashes",
|
|
"--emit-index-url",
|
|
"--emit-find-links",
|
|
"--no-strip-extras",
|
|
"--index-strategy=unsafe-best-match",
|
|
"--upgrade",
|
|
],
|
|
python_platform = "x86_64-unknown-linux-gnu",
|
|
requirements_in = "requirements.in",
|
|
requirements_txt = "requirements.lock.txt",
|
|
)
|
|
|
|
py_console_script_binary(
|
|
name = "neuronx-cc",
|
|
binary_rule = py_binary_with_script,
|
|
pkg = "@neuron_py_deps//neuronx_cc",
|
|
visibility = ["//visibility:public"],
|
|
)
|
|
|
|
py_binary_with_script(
|
|
name = "libneuronxla",
|
|
srcs = ["empty.py"],
|
|
main = "empty.py",
|
|
deps = [
|
|
":neuronx-cc",
|
|
"@neuron_py_deps//libneuronxla",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "libpython",
|
|
hdrs = ["libpython.h"],
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
"@rules_python//python/cc:current_py_cc_headers",
|
|
"@rules_python//python/cc:current_py_cc_libs",
|
|
],
|
|
)
|
|
|
|
compile_pip_requirements(
|
|
name = "requirements",
|
|
src = "requirements.in",
|
|
py_binary = py_binary_with_script,
|
|
requirements_txt = "requirements.lock.txt",
|
|
)
|
|
|
|
runfiles_to_default(
|
|
name = "libneuronxla_files",
|
|
deps = ["@neuron_py_deps//libneuronxla:pkg"],
|
|
)
|
|
|
|
select_file(
|
|
name = "libneuronpjrt_so",
|
|
srcs = ":libneuronxla_files",
|
|
subpath = "site-packages/libneuronxla/libneuronpjrt.so",
|
|
)
|
|
|
|
cc_import(
|
|
name = "libpjrt_neuron_",
|
|
add_needed = [
|
|
"libpython3.11.so.1.0",
|
|
"libzmlxneuron.so.0",
|
|
],
|
|
data = [":libneuronxla"],
|
|
shared_library = ":libneuronpjrt_so",
|
|
soname = "libpjrt_neuron.so",
|
|
deps = [
|
|
":zmlxneuron",
|
|
"@aws-neuronx-runtime-lib",
|
|
"@rules_python//python/cc:current_py_cc_libs",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "zmlxneuron_lib",
|
|
srcs = ["zmlxneuron.c"],
|
|
linkopts = ["-ldl"],
|
|
)
|
|
|
|
cc_shared_library(
|
|
name = "zmlxneuron_",
|
|
shared_lib_name = "libzmlxneuron.so.0",
|
|
deps = [":zmlxneuron_lib"],
|
|
)
|
|
|
|
cc_import(
|
|
name = "zmlxneuron",
|
|
shared_library = ":zmlxneuron_",
|
|
)
|
|
|
|
alias(
|
|
name = "bootstrap",
|
|
actual = "@rules_python//python/config_settings:bootstrap_impl",
|
|
)
|
|
|
|
cc_library(
|
|
name = "empty",
|
|
)
|
|
|
|
cc_library(
|
|
name = "libpjrt_neuron",
|
|
hdrs = ["libpjrt_neuron.h"],
|
|
defines = ["ZML_RUNTIME_NEURON"],
|
|
deps = [":libpjrt_neuron_"],
|
|
)
|
|
|
|
zig_library(
|
|
name = "neuron",
|
|
import_name = "runtimes/neuron",
|
|
main = "neuron.zig",
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
"//pjrt",
|
|
] + select({
|
|
"//runtimes:neuron.enabled": [
|
|
":libneuronxla_pyenv",
|
|
":libpjrt_neuron",
|
|
":libpython",
|
|
"//async",
|
|
"@rules_zig//zig/runfiles",
|
|
],
|
|
"//conditions:default": [":empty"],
|
|
}),
|
|
)
|
|
|
|
filegroup(
|
|
name = "layers",
|
|
srcs = [
|
|
"@apt_neuron//bash/amd64",
|
|
"@apt_neuron//libxml2/amd64",
|
|
"@apt_neuron//zlib1g/amd64",
|
|
],
|
|
visibility = ["//visibility:public"],
|
|
)
|