From 1a2b862ec217bf46a16792dae15ce5a01ea3995d Mon Sep 17 00:00:00 2001 From: Tarry Singh Date: Mon, 19 May 2025 17:35:33 +0000 Subject: [PATCH] Add sandbox neuron dependencies: define a trampoline PJRT, create an empty repository for distroless deps, and update Bazel build files and Zig/C sources accordingly. --- MODULE.bazel | 2 +- MODULE.bazel.lock | 46 +++++++++-- bazel/simple_repository.bzl | 23 ++++++ runtimes/common/packages.bzl | 11 +++ runtimes/neuron/BUILD.bazel | 79 ++++++++---------- runtimes/neuron/libpjrt_neuron.BUILD.bazel | 96 ++++++++++++++++++++++ runtimes/neuron/libpjrt_neuron.zig | 39 +++++++++ runtimes/neuron/neuron.bzl | 85 +++++++++++++------ runtimes/neuron/neuron.zig | 35 ++++++-- runtimes/neuron/packages.lock.json | 89 ++++---------------- runtimes/neuron/packages.yaml | 6 +- runtimes/neuron/requirements.lock.txt | 49 +++++------ runtimes/neuron/zmlxneuron.c | 5 +- 13 files changed, 375 insertions(+), 190 deletions(-) create mode 100644 bazel/simple_repository.bzl create mode 100644 runtimes/neuron/libpjrt_neuron.BUILD.bazel create mode 100644 runtimes/neuron/libpjrt_neuron.zig diff --git a/MODULE.bazel b/MODULE.bazel index 401868a..ce4383a 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -92,7 +92,7 @@ tpu = use_extension("//runtimes/tpu:tpu.bzl", "tpu_packages") use_repo(tpu, "libpjrt_tpu") neuron = use_extension("//runtimes/neuron:neuron.bzl", "neuron_packages") -use_repo(neuron, "aws-neuronx-collectives", "aws-neuronx-runtime-lib") +use_repo(neuron, "libpjrt_neuron") zls = use_extension("//third_party/zls:zls.bzl", "repo") use_repo(zls, "zls_aarch64-macos", "zls_x86_64-linux", "zls_x86_64-macos") diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 1c74b30..cca6007 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1033,7 +1033,7 @@ "bzlTransitiveDigest": "fOwLlf5Rfxv4xkFCZwPlkSPUuZ9lAhMGf0SN43pc0xg=", "usagesDigest": "3+0h/pprjG6w7aHgvrT2H3LdTTVQQ79Mz0lATHIw+/I=", "recordedFileInputs": { - "@@//runtimes/neuron/requirements.lock.txt": "a5e5353b149b33dc067ce27ef67896d8ca3518e37a344e215e6604f98ea6d0fc", + "@@//runtimes/neuron/requirements.lock.txt": "fa0b5b967162e303415ee1dd7d2c1a448068935663ba5ac5b5f72a515fe3f6f6", "@@protobuf+//python/requirements.txt": "2bb341a0ba317fb8dc9ac79817b9b2d8f6a4d0c43e3fb652f0cab584eed420a4", "@@rules_fuzzing+//fuzzing/requirements.txt": "d4e91a1b3eca06d6e2877f67d3d8be498cb3ed6d28b9c14439b6f210ec621470", "@@rules_python+//tools/publish/requirements_darwin.txt": "6d10732737e5504291360b78fbf2b82abe39a5a021bd937f987f8433519b0f2c", @@ -1069,7 +1069,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "boto3==1.38.25 --hash=sha256:2f2cd517dd31d33ace0eefe567dc903fdf74221513e32f1e9445bdfac7554db7 --hash=sha256:85c1556a110896f68de8573a9b4757c81071448dbf6ffc1074941bfc8a43195e" + "requirement": "boto3==1.39.9 --hash=sha256:5bc85e9fdec4e21ef5ca2c22b4d51a3e32b53f3da36ce51f5a3ea4dbde07b132 --hash=sha256:e3d3a6b617e1575e7ec854c820a882ab2e189a0421e74dc0dca2c9e13d4370a5" } }, "neuron_py_deps_311_botocore": { @@ -1095,7 +1095,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "botocore==1.38.25 --hash=sha256:5a960bd990a11cdb78865e908a58ed712d87d9b419f55ad21c99d7d21f0d6726 --hash=sha256:8c73e97d9662a6c92be33dab66cd1e2b59797154c7ec379ce3bb5d6779d9d78c" + "requirement": "botocore==1.39.9 --hash=sha256:02f141c2849e4589a79feea245ce4ecc478d48b7865572445af8aae3b041772d --hash=sha256:a9691cbe03a3bc8b2720b3c36e5c5a2eecace6acd72bfb1107f00e75edaec4f3" } }, "neuron_py_deps_311_certifi": { @@ -1121,7 +1121,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "certifi==2025.4.26 --hash=sha256:0a816057ea3cdefcef70270d2c515e4506bbc954f417fa5ade2021213bb8f0c6 --hash=sha256:30350364dfe371162649852c63336a15c70c6510c2ad5015b21c2345311805f3" + "requirement": "certifi==2025.7.14 --hash=sha256:6b31f564a415d79ee77df69d757bb49a5bb53bd9f756cbbe24394ffd6fc1f4b2 --hash=sha256:8ea99dbdfaaf2ba2f9bac77b9249ef62ec5218e7c2b2e903378ed5fccf765995" } }, "neuron_py_deps_311_charset_normalizer": { @@ -1277,7 +1277,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "libneuronxla==2.2.3493.0+78c3e78c --hash=sha256:9ba54e5d481afc25ae83e964a7c076543991ad80b88839cd500abcac130595e8" + "requirement": "libneuronxla==2.2.4410.0+835a67fb --hash=sha256:37e5f08482ef1a9a844c2894de6d91c03400a092764e54ab67612456be439f16" } }, "neuron_py_deps_311_lockfile": { @@ -1381,7 +1381,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "neuronx-cc==2.18.121.0+9e31e41a --hash=sha256:b4c0959594af716aff94698e15cc3997cc39a66c853c25a1bbf5eb7bf71fd7ec" + "requirement": "neuronx-cc==2.19.8089.0+8ab9f450 --hash=sha256:fbb05d8a0da8f54954c03ff2bea659f627309110c00efd54f523d877ddd6b88d" } }, "neuron_py_deps_311_numpy": { @@ -1563,7 +1563,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "requests==2.31.0 --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f --hash=sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1" + "requirement": "requests==2.32.4 --hash=sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c --hash=sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422" } }, "neuron_py_deps_311_requests_unixsocket": { @@ -1615,7 +1615,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "s3transfer==0.13.0 --hash=sha256:0148ef34d6dd964d0d8cf4311b2b21c474693e57c2e069ec708ce043d2b527be --hash=sha256:f5e6db74eb7776a37208001113ea7aa97695368242b364d73e91c981ac522177" + "requirement": "s3transfer==0.13.1 --hash=sha256:a981aa7429be23fe6dfc13e80e4020057cbab622b08c0315288758d67cabc724 --hash=sha256:c3fdba22ba1bd367922f27ec8032d6a1cf5f10c934fb5d68cf60fd5a23d936cf" } }, "neuron_py_deps_311_scipy": { @@ -1696,6 +1696,32 @@ "requirement": "tqdm==4.67.1 --hash=sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2 --hash=sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2" } }, + "neuron_py_deps_311_typing_extensions": { + "repoRuleId": "@@rules_python+//python/private/pypi:whl_library.bzl%whl_library", + "attributes": { + "dep_template": "@neuron_py_deps//{name}:{target}", + "download_only": true, + "experimental_target_platforms": [ + "cp311_linux_x86_64" + ], + "extra_pip_args": [ + "--index-url", + "https://pypi.org/simple", + "--extra-index-url", + "https://pip.repos.neuron.amazonaws.com", + "--find-links", + "https://mirror.zml.ai/pypi/aws-neuronx-runtime-discovery/index.html", + "--abi=cp311", + "--implementation=cp", + "--python-version=311", + "--platform=linux_x86_64", + "--platform=manylinux2014_x86_64" + ], + "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", + "repo": "neuron_py_deps_311", + "requirement": "typing-extensions==4.14.1 --hash=sha256:38b39f4aeeab64884ce9f74c94263ef78f3c22467c8724005483154c26648d36 --hash=sha256:d1e1e3b58374dc93031d6eda2420a48ea44a36c2b4766a4fdeb3710755731d76" + } + }, "neuron_py_deps_311_urllib3": { "repoRuleId": "@@rules_python+//python/private/pypi:whl_library.bzl%whl_library", "attributes": { @@ -1719,7 +1745,7 @@ ], "python_interpreter_target": "@@rules_python++python+python_3_11_host//:python", "repo": "neuron_py_deps_311", - "requirement": "urllib3==2.4.0 --hash=sha256:414bc6535b787febd7567804cc015fee39daab8ad86268f1310a9250697de466 --hash=sha256:4e16665048960a0900c702d4a66415956a584919c03361cac9f1df5c5dd7e813" + "requirement": "urllib3==2.5.0 --hash=sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760 --hash=sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc" } }, "pip_deps_310_numpy": { @@ -4148,6 +4174,7 @@ "scipy": "{\"neuron_py_deps_311_scipy\":[{\"version\":\"3.11\"}]}", "six": "{\"neuron_py_deps_311_six\":[{\"version\":\"3.11\"}]}", "tqdm": "{\"neuron_py_deps_311_tqdm\":[{\"version\":\"3.11\"}]}", + "typing_extensions": "{\"neuron_py_deps_311_typing_extensions\":[{\"version\":\"3.11\"}]}", "urllib3": "{\"neuron_py_deps_311_urllib3\":[{\"version\":\"3.11\"}]}" }, "packages": [ @@ -4176,6 +4203,7 @@ "scipy", "six", "tqdm", + "typing_extensions", "urllib3" ], "groups": {} diff --git a/bazel/simple_repository.bzl b/bazel/simple_repository.bzl new file mode 100644 index 0000000..9fd0713 --- /dev/null +++ b/bazel/simple_repository.bzl @@ -0,0 +1,23 @@ +_UNSET = "_UNSET" + +def _simple_repository_impl(rctx): + if (rctx.attr.build_file == None) == (rctx.attr.build_file_content == _UNSET): + fail("exactly one of `build_file` and `build_file_content` must be specified") + + if rctx.attr.build_file != None: + # Remove any existing BUILD.bazel in the repository to ensure + # the symlink to the defined build_file doesn't fail. + rctx.delete("BUILD.bazel") + rctx.symlink(rctx.attr.build_file, "BUILD.bazel") + else: + rctx.file("BUILD.bazel", rctx.attr.build_file_content) + +simple_repository = repository_rule( + implementation = _simple_repository_impl, + attrs = { + "build_file": attr.label(allow_single_file = True), + "build_file_content": attr.string(default = _UNSET), + }, + doc = "Makes an empty repository from just one BUILD.bazel file.", + local = True, +) diff --git a/runtimes/common/packages.bzl b/runtimes/common/packages.bzl index 8e3c1c4..abb2ab1 100644 --- a/runtimes/common/packages.bzl +++ b/runtimes/common/packages.bzl @@ -13,6 +13,16 @@ def _cc_import(**kwargs): def _cc_library(**kwargs): return """cc_library({})""".format(_kwargs(**kwargs)) +def _cc_library_hdrs_glob(name, hdrs_glob, deps = [], **kwargs): + return """\ +cc_library( + name = "{name}", + hdrs = glob({hdrs_glob}), + deps = {deps}, + {kwargs} +) +""".format(name = name, hdrs_glob = repr(hdrs_glob), deps = repr(deps), kwargs = _kwargs(**kwargs)) + def _cc_import_glob_hdrs(name, hdrs_glob, shared_library, deps = [], **kwargs): return """\ cc_import( @@ -48,6 +58,7 @@ packages = struct( cc_import = _cc_import, cc_import_glob_hdrs = _cc_import_glob_hdrs, cc_library = _cc_library, + cc_library_hdrs_glob = _cc_library_hdrs_glob, filegroup = _filegroup, load_ = _load, patchelf = _patchelf, diff --git a/runtimes/neuron/BUILD.bazel b/runtimes/neuron/BUILD.bazel index 01a3684..b65d7a3 100644 --- a/runtimes/neuron/BUILD.bazel +++ b/runtimes/neuron/BUILD.bazel @@ -1,12 +1,8 @@ -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("@rules_zig//zig:defs.bzl", "BINARY_KIND", "zig_binary", "zig_library") load(":neuron.bzl", "py_binary_with_script") load(":pyenv.bzl", "pyenv_zig") @@ -20,6 +16,25 @@ zig_library( main = ":libneuronxla_pyenv_zig", ) +# A proxy PJRT Plugin that loads the Neuron PJRT Plugin +# and returns the instance from nested GetPjrtApi. +# +# Additionally, it provides a way to load implicit transitive dependencies +# of neuronx-cc (see add_needed of the patchelf target below). +# +# TODO(cerisier): Use a zig_cc_shared_library instead. +zig_binary( + name = "libpjrt_neuron_proxy", + kind = BINARY_KIND.shared_lib, + main = "libpjrt_neuron.zig", + copts = ["-lc"], + deps = [ + "//stdx", + "@rules_zig//zig/runfiles", + ], + visibility = ["@libpjrt_neuron//:__subpackages__"], +) + pip_compile( name = "update_requirements", args = [ @@ -75,48 +90,11 @@ compile_pip_requirements( ], ) -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_", + visibility = ["@libpjrt_neuron//:__subpackages__"], ) alias( @@ -132,7 +110,9 @@ cc_library( name = "libpjrt_neuron", hdrs = ["libpjrt_neuron.h"], defines = ["ZML_RUNTIME_NEURON"], - deps = [":libpjrt_neuron_"], + deps = [ + "@libpjrt_neuron", + ], ) zig_library( @@ -140,14 +120,21 @@ zig_library( import_name = "runtimes/neuron", main = "neuron.zig", visibility = ["//visibility:public"], + data = select({ + "//runtimes:neuron.enabled": [ + ":libneuronxla", + ], + "//conditions:default": [], + }), deps = [ "//pjrt", ] + select({ "//runtimes:neuron.enabled": [ - ":libneuronxla_pyenv", ":libpjrt_neuron", + ":libneuronxla_pyenv", ":libpython", "//async", + "//stdx", "@rules_zig//zig/runfiles", ], "//conditions:default": [":empty"], @@ -158,8 +145,6 @@ filegroup( name = "layers", srcs = [ "@apt_neuron//bash/amd64", - "@apt_neuron//libxml2/amd64", - "@apt_neuron//zlib1g/amd64", ], visibility = ["//visibility:public"], ) diff --git a/runtimes/neuron/libpjrt_neuron.BUILD.bazel b/runtimes/neuron/libpjrt_neuron.BUILD.bazel new file mode 100644 index 0000000..5df16d6 --- /dev/null +++ b/runtimes/neuron/libpjrt_neuron.BUILD.bazel @@ -0,0 +1,96 @@ +load("@aspect_bazel_lib//lib:copy_to_directory.bzl", "copy_to_directory") +load("@zml//bazel:patchelf.bzl", "patchelf") +load("@bazel_skylib//rules:copy_file.bzl", "copy_file") +load("@bazel_skylib//rules:select_file.bzl", "select_file") +load("@rules_cc//cc:cc_shared_library.bzl", "cc_shared_library") +load("@zml//bazel:runfiles.bzl", "runfiles_to_default") + +cc_shared_library( + name = "zmlxneuron_so_", + shared_lib_name = "libzmlxneuron.so.0", + deps = ["@zml//runtimes/neuron:zmlxneuron_lib"], +) + +copy_file( + name = "zmlxneuron_so", + src = ":zmlxneuron_so_", + out = "lib/libzmlxneuron.so.0", +) + +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", +) + +patchelf( + name = "libneuronpjrt.patchelf", + add_needed = [ + "libpython3.11.so.1.0", + "libzmlxneuron.so.0", + "libnccom.so.2", + ], + set_rpath = '$ORIGIN', + shared_library = ":libneuronpjrt_so", + soname = "libneuronpjrt.so", +) + +patchelf( + name = "libpjrt_neuron_proxy.patchelf", + set_rpath = '$ORIGIN', + add_needed = [ + "libz.so.1", + "libgomp.so.1", + ], + shared_library = "@zml//runtimes/neuron:libpjrt_neuron_proxy", + soname = "libpjrt_neuron.so", +) + +copy_to_directory( + name = "sandbox", + srcs = [ + ":zmlxneuron_so", + ":libneuronpjrt.patchelf", + ":libpjrt_neuron_proxy.patchelf", + "@aws-neuronx-runtime-lib//:libnrt.patchelf", + "@aws-neuronx-runtime-lib//:libncfw.patchelf", + "@aws-neuronx-collectives//:libnccom", + "@zlib1g", + "@libgomp1", + ], + replace_prefixes = { + "libneuronpjrt.patchelf": "lib", + "libpjrt_neuron_proxy.patchelf": "lib", + "libnrt.patchelf": "lib", + "libncfw.patchelf": "lib", + "lib/x86_64-linux-gnu": "lib", + "usr/lib/x86_64-linux-gnu": "lib", + "opt/neuron": "lib", + }, + add_directory_to_runfiles = True, + include_external_repositories = ["**"], +) + +cc_library( + name = "libpjrt_neuron", + data = [":sandbox"], + deps = [ + "@aws-neuronx-runtime-lib//:libnrt_headers", + ], + linkopts = [ + # Defer function call resolution until the function is called + # (lazy loading) rather than at load time. + # + # This is required because we want to let downstream use weak NRT symbols. + # + # We force it here because -z,now (which resolve all symbols at load time), + # is the default in most bazel CC toolchains as well as in certain linkers. + "-Wl,-z,lazy", + ], + visibility = ["@zml//runtimes/neuron:__subpackages__"], +) diff --git a/runtimes/neuron/libpjrt_neuron.zig b/runtimes/neuron/libpjrt_neuron.zig new file mode 100644 index 0000000..6485ecd --- /dev/null +++ b/runtimes/neuron/libpjrt_neuron.zig @@ -0,0 +1,39 @@ +const std = @import("std"); +const runfiles = @import("runfiles"); +const bazel_builtin = @import("bazel_builtin"); +const stdx = @import("stdx"); + +const log = std.log.scoped(.@"zml/runtimes/neuron"); + +pub export fn GetPjrtApi() *anyopaque { + var arena = std.heap.ArenaAllocator.init(std.heap.c_allocator); + defer arena.deinit(); + + var r_ = runfiles.Runfiles.create(.{ .allocator = arena.allocator() }) catch |err| { + stdx.debug.panic("Unable to find runfiles: {}", .{err}); + } orelse stdx.debug.panic("Runfiles not availeabwewefle", .{}); + + const source_repo = bazel_builtin.current_repository; + const r = r_.withSourceRepo(source_repo); + + var path_buf: [std.fs.max_path_bytes]u8 = undefined; + const sandbox_path = r.rlocation("libpjrt_neuron/sandbox", &path_buf) catch |err| { + stdx.debug.panic("Failed to find sandbox path for NEURON runtime: {}", .{err}); + } orelse stdx.debug.panic("No NEURON sandbox path found", .{}); + + var lib_path_buf: [std.fs.max_path_bytes]u8 = undefined; + const library = stdx.fs.path.bufJoinZ(&lib_path_buf, &.{ sandbox_path, "lib", "libneuronpjrt.so" }) catch unreachable; + + var lib: std.DynLib = blk: { + const handle = std.c.dlopen(library, .{ .LAZY = true, .GLOBAL = true, .NODELETE = true }) orelse { + stdx.debug.panic("Unable to dlopen plugin: {s}", .{library}); + }; + break :blk .{ .inner = .{ .handle = handle } }; + }; + + const sym = lib.lookup(*const fn () callconv(.C) *anyopaque, "GetPjrtApi") orelse { + stdx.debug.panic("Unable to find symbol GetPjrtApi in plugin: {s}", .{library}); + }; + + return sym(); +} diff --git a/runtimes/neuron/neuron.bzl b/runtimes/neuron/neuron.bzl index 4466ca9..89996fa 100644 --- a/runtimes/neuron/neuron.bzl +++ b/runtimes/neuron/neuron.bzl @@ -2,39 +2,55 @@ load("@python_versions//3.11:defs.bzl", _py_binary = "py_binary") load("@rules_python//python:defs.bzl", "PyInfo") load("@with_cfg.bzl", "with_cfg") load("//bazel:http_deb_archive.bzl", "http_deb_archive") +load("//bazel:simple_repository.bzl", "simple_repository") load("//runtimes/common:packages.bzl", "packages") BASE_URL = "https://apt.repos.neuron.amazonaws.com" STRIP_PREFIX = "opt/aws/neuron" -BUILD_FILE_PRELUDE = """\ -load("@zml//bazel:cc_import.bzl", "cc_import") +_BUILD_FILE_PRELUDE = """\ +package(default_visibility = ["//visibility:public"]) """ -_PACKAGES = { - "aws-neuronx-runtime-lib": packages.cc_import_glob_hdrs( - name = "aws-neuronx-runtime-lib", - hdrs_glob = [ - "include/ndl/**/*.h", - "include/nrt/**/*.h", - ], - includes = ["include"], - shared_library = "lib/libnrt.so.1", - rename_dynamic_symbols = { - "dlopen": "zmlxneuron_dlopen", - }, - visibility = ["@zml//runtimes/neuron:__subpackages__"], - deps = ["@aws-neuronx-collectives//:libnccom"], - ), - "aws-neuronx-collectives": "\n".join([ - packages.cc_import( - name = "libnccom", - shared_library = "lib/libnccom.so.2", - visibility = ["@aws-neuronx-runtime-lib//:__subpackages__"], +_UBUNTU_PACKAGES = { + "zlib1g": packages.filegroup(name = "zlib1g", srcs = ["lib/x86_64-linux-gnu/libz.so.1"]), + "libgomp1": packages.filegroup(name = "libgomp1", srcs = ["usr/lib/x86_64-linux-gnu/libgomp.so.1"]), +} + +_NEURON_PACKAGES = { + "aws-neuronx-runtime-lib": "\n".join([ + packages.load_("@zml//bazel:patchelf.bzl", "patchelf"), + packages.cc_library_hdrs_glob( + name = "libnrt_headers", + hdrs_glob = [ + "include/ndl/**/*.h", + "include/nrt/**/*.h", + ], + includes = ["include"], + visibility = ["//visibility:public"], ), - packages.cc_import( - name = "libnccom-net", - shared_library = "lib/libnccom-net.so.0", + packages.patchelf( + name = "libnrt.patchelf", + shared_library = "lib/libnrt.so.1", + set_rpath = '$ORIGIN', + add_needed = [ + # readelf -d ./opt/aws/neuron/libl/libncfw.so + "libncfw.so.2", + ], + rename_dynamic_symbols = { + "dlopen": "zmlxneuron_dlopen", + }, + ), + packages.patchelf( + name = "libncfw.patchelf", + shared_library = "lib/libncfw.so", + soname = "libncfw.so.2", + ), + ]), + "aws-neuronx-collectives": "\n".join([ + packages.filegroup( + name = "libnccom", + srcs = ["lib/libnccom.so.2"], ), ]), } @@ -43,14 +59,29 @@ def _neuron_impl(mctx): loaded_packages = packages.read(mctx, [ "@zml//runtimes/neuron:packages.lock.json", ]) - for pkg_name, build_file_content in _PACKAGES.items(): + + simple_repository( + name = "libpjrt_neuron", + build_file = ":libpjrt_neuron.BUILD.bazel", + ) + + for pkg_name, build_file_content in _UBUNTU_PACKAGES.items(): + pkg = loaded_packages[pkg_name] + http_deb_archive( + name = pkg_name, + urls = pkg["urls"], + sha256 = pkg["sha256"], + build_file_content = _BUILD_FILE_PRELUDE + build_file_content, + ) + + for pkg_name, build_file_content in _NEURON_PACKAGES.items(): pkg = loaded_packages[pkg_name] http_deb_archive( name = pkg_name, urls = pkg["urls"], sha256 = pkg["sha256"], strip_prefix = STRIP_PREFIX, - build_file_content = BUILD_FILE_PRELUDE + build_file_content, + build_file_content = _BUILD_FILE_PRELUDE + build_file_content, ) return mctx.extension_metadata( diff --git a/runtimes/neuron/neuron.zig b/runtimes/neuron/neuron.zig index 9fdbcc4..fc97342 100644 --- a/runtimes/neuron/neuron.zig +++ b/runtimes/neuron/neuron.zig @@ -7,6 +7,9 @@ const c = @import("c"); const libneuronxla_pyenv = @import("libneuronxla_pyenv"); const pjrt = @import("pjrt"); const runfiles = @import("runfiles"); +const stdx = @import("stdx"); + +const log = std.log.scoped(.@"zml/runtime/neuron"); pub fn isEnabled() bool { return @hasDecl(c, "ZML_RUNTIME_NEURON"); @@ -44,11 +47,7 @@ fn pyErrorOrExit(status: c.PyStatus) void { } } -fn initialize() !void { - var arena = std.heap.ArenaAllocator.init(std.heap.c_allocator); - const allocator = arena.allocator(); - defer arena.deinit(); - +fn initialize(allocator: std.mem.Allocator, r_: *runfiles.Runfiles) !void { { var preconfig: c.PyPreConfig = undefined; c.PyPreConfig_InitIsolatedConfig(&preconfig); @@ -60,7 +59,6 @@ fn initialize() !void { c.PyConfig_InitIsolatedConfig(&config); defer c.PyConfig_Clear(&config); - var r_ = try runfiles.Runfiles.create(.{ .allocator = allocator }) orelse return error.Unavailable; const r = r_.withSourceRepo(bazel_builtin.current_repository); var buf: [std.fs.max_path_bytes]u8 = undefined; @@ -135,7 +133,28 @@ pub fn load() !*const pjrt.Api { return error.Unavailable; } + var arena = std.heap.ArenaAllocator.init(std.heap.c_allocator); + defer arena.deinit(); + + var r_ = try runfiles.Runfiles.create(.{ .allocator = arena.allocator() }) orelse { + stdx.debug.panic("Unable to find runfiles", .{}); + }; + + const source_repo = bazel_builtin.current_repository; + const r = r_.withSourceRepo(source_repo); + + var path_buf: [std.fs.max_path_bytes]u8 = undefined; + const sandbox_path = try r.rlocation("libpjrt_neuron/sandbox", &path_buf) orelse { + log.err("Failed to find sandbox path for NEURON runtime", .{}); + return error.FileNotFound; + }; + setNeuronCCFlags(); - try initialize(); - return try asynk.callBlocking(pjrt.Api.loadFrom, .{"libpjrt_neuron.so"}); + try initialize(arena.allocator(), &r_); + + return blk: { + var lib_path_buf: [std.fs.max_path_bytes]u8 = undefined; + const path = try stdx.fs.path.bufJoinZ(&lib_path_buf, &.{ sandbox_path, "lib", "libpjrt_neuron.so" }); + break :blk asynk.callBlocking(pjrt.Api.loadFrom, .{path}); + }; } diff --git a/runtimes/neuron/packages.lock.json b/runtimes/neuron/packages.lock.json index 2ca0809..ff484ae 100755 --- a/runtimes/neuron/packages.lock.json +++ b/runtimes/neuron/packages.lock.json @@ -143,22 +143,17 @@ "version": "12.2.0-14+deb12u1" } ], - "key": "zlib1g_1-1.2.13.dfsg-1_amd64", - "name": "zlib1g", - "sha256": "d7dd1d1411fedf27f5e27650a6eff20ef294077b568f4c8c5e51466dc7c08ce4", + "key": "libgomp1_12.2.0-14-p-deb12u1_amd64", + "name": "libgomp1", + "sha256": "48fec46bda7f5b1638b9e959889bfbc20491247d402d120bb152687eb48143d7", "urls": [ - "https://snapshot-cloudflare.debian.org/archive/debian/20250529T205323Z/pool/main/z/zlib/zlib1g_1.2.13.dfsg-1_amd64.deb" + "https://snapshot-cloudflare.debian.org/archive/debian/20250529T205323Z/pool/main/g/gcc-12/libgomp1_12.2.0-14+deb12u1_amd64.deb" ], - "version": "1:1.2.13.dfsg-1" + "version": "12.2.0-14+deb12u1" }, { "arch": "amd64", "dependencies": [ - { - "key": "zlib1g_1-1.2.13.dfsg-1_amd64", - "name": "zlib1g", - "version": "1:1.2.13.dfsg-1" - }, { "key": "libc6_2.36-9-p-deb12u10_amd64", "name": "libc6", @@ -173,74 +168,26 @@ "key": "gcc-12-base_12.2.0-14-p-deb12u1_amd64", "name": "gcc-12-base", "version": "12.2.0-14+deb12u1" - }, - { - "key": "liblzma5_5.4.1-1_amd64", - "name": "liblzma5", - "version": "5.4.1-1" - }, - { - "key": "libicu72_72.1-3_amd64", - "name": "libicu72", - "version": "72.1-3" - }, - { - "key": "libstdc-p--p-6_12.2.0-14-p-deb12u1_amd64", - "name": "libstdc++6", - "version": "12.2.0-14+deb12u1" } ], - "key": "libxml2_2.9.14-p-dfsg-1.3_deb12u1_amd64", - "name": "libxml2", - "sha256": "35b76cb7038fc1c940204a4f05f33ffb79d027353ce469397d9adcf8f9b3e1a7", + "key": "zlib1g_1-1.2.13.dfsg-1_amd64", + "name": "zlib1g", + "sha256": "d7dd1d1411fedf27f5e27650a6eff20ef294077b568f4c8c5e51466dc7c08ce4", "urls": [ - "https://snapshot-cloudflare.debian.org/archive/debian/20250529T205323Z/pool/main/libx/libxml2/libxml2_2.9.14+dfsg-1.3~deb12u1_amd64.deb" + "https://snapshot-cloudflare.debian.org/archive/debian/20250529T205323Z/pool/main/z/zlib/zlib1g_1.2.13.dfsg-1_amd64.deb" ], - "version": "2.9.14+dfsg-1.3~deb12u1" + "version": "1:1.2.13.dfsg-1" }, { "arch": "amd64", "dependencies": [], - "key": "liblzma5_5.4.1-1_amd64", - "name": "liblzma5", - "sha256": "d321b9502b16aac534e1c691afbe3dc5e125e5091aa35bea026c59b25ebe82e7", - "urls": [ - "https://snapshot-cloudflare.debian.org/archive/debian/20250529T205323Z/pool/main/x/xz-utils/liblzma5_5.4.1-1_amd64.deb" - ], - "version": "5.4.1-1" - }, - { - "arch": "amd64", - "dependencies": [], - "key": "libicu72_72.1-3_amd64", - "name": "libicu72", - "sha256": "e239c1c9f52bee0ff627f291552d63691b765ec7c5cdf6de7c7ae4dec0275857", - "urls": [ - "https://snapshot-cloudflare.debian.org/archive/debian/20250529T205323Z/pool/main/i/icu/libicu72_72.1-3_amd64.deb" - ], - "version": "72.1-3" - }, - { - "arch": "amd64", - "dependencies": [], - "key": "libstdc-p--p-6_12.2.0-14-p-deb12u1_amd64", - "name": "libstdc++6", - "sha256": "5cd3171216d4ab0fc911cfe9c35509bf2dd8f47761c43b7f6a4296701551a24d", - "urls": [ - "https://snapshot-cloudflare.debian.org/archive/debian/20250529T205323Z/pool/main/g/gcc-12/libstdc++6_12.2.0-14+deb12u1_amd64.deb" - ], - "version": "12.2.0-14+deb12u1" - }, - { - "arch": "amd64", - "dependencies": [], - "key": "aws-neuronx-runtime-lib_2.25.57.0-166c7a468_amd64", + "key": "aws-neuronx-runtime-lib_2.26.42.0-2ff3b5c7d_amd64", "name": "aws-neuronx-runtime-lib", - "sha256": "81613a0b86fc704acb5a4d1096a5bf2cce3aead04aeb5063d14a95195ca8ca39", + "sha256": "57f6fd8e850f4054c2d229cd41a83fe1489ec9d9ee55bedf33a14cd6e8beed16", "urls": [ - "https://apt.repos.neuron.amazonaws.com/pool/main/a/aws-neuronx-runtime-lib/aws-neuronx-runtime-lib_2.25.57.0-166c7a468_amd64.deb" + "https://apt.repos.neuron.amazonaws.com/pool/main/a/aws-neuronx-runtime-lib/aws-neuronx-runtime-lib_2.26.42.0-2ff3b5c7d_amd64.deb" ], - "version": "2.25.57.0-166c7a468" + "version": "2.26.42.0-2ff3b5c7d" }, { "arch": "amd64", @@ -386,13 +333,13 @@ "version": "252.36-1~deb12u1" } ], - "key": "aws-neuronx-collectives_2.25.65.0-9858ac9a1_amd64", + "key": "aws-neuronx-collectives_2.26.43.0-47cc904ea_amd64", "name": "aws-neuronx-collectives", - "sha256": "41e2eb08433556b243bae12ea61a7417068374b674b5a4e3e88455d28b6ba345", + "sha256": "4886f5aeeea21e4d502d445654a838190e2174a52ffb757da2cc3a8bf09be3f8", "urls": [ - "https://apt.repos.neuron.amazonaws.com/pool/main/a/aws-neuronx-collectives/aws-neuronx-collectives_2.25.65.0-9858ac9a1_amd64.deb" + "https://apt.repos.neuron.amazonaws.com/pool/main/a/aws-neuronx-collectives/aws-neuronx-collectives_2.26.43.0-47cc904ea_amd64.deb" ], - "version": "2.25.65.0-9858ac9a1" + "version": "2.26.43.0-47cc904ea" }, { "arch": "amd64", diff --git a/runtimes/neuron/packages.yaml b/runtimes/neuron/packages.yaml index c0d576a..6c3de28 100644 --- a/runtimes/neuron/packages.yaml +++ b/runtimes/neuron/packages.yaml @@ -18,8 +18,8 @@ archs: - "amd64" packages: - - "bash" - - "zlib1g" - - "libxml2" + - "bash" # needed to launch neuronx-cc + - "libgomp1" # implicit dep of neuronx-cc + - "zlib1g" # implicit dep of libpython - "aws-neuronx-runtime-lib" - "aws-neuronx-collectives" diff --git a/runtimes/neuron/requirements.lock.txt b/runtimes/neuron/requirements.lock.txt index ab1e288..c5994c0 100644 --- a/runtimes/neuron/requirements.lock.txt +++ b/runtimes/neuron/requirements.lock.txt @@ -4,20 +4,20 @@ --extra-index-url https://pip.repos.neuron.amazonaws.com --find-links https://mirror.zml.ai/pypi/aws-neuronx-runtime-discovery/index.html -boto3==1.38.25 \ - --hash=sha256:2f2cd517dd31d33ace0eefe567dc903fdf74221513e32f1e9445bdfac7554db7 \ - --hash=sha256:85c1556a110896f68de8573a9b4757c81071448dbf6ffc1074941bfc8a43195e +boto3==1.39.9 \ + --hash=sha256:5bc85e9fdec4e21ef5ca2c22b4d51a3e32b53f3da36ce51f5a3ea4dbde07b132 \ + --hash=sha256:e3d3a6b617e1575e7ec854c820a882ab2e189a0421e74dc0dca2c9e13d4370a5 # via libneuronxla -botocore==1.38.25 \ - --hash=sha256:5a960bd990a11cdb78865e908a58ed712d87d9b419f55ad21c99d7d21f0d6726 \ - --hash=sha256:8c73e97d9662a6c92be33dab66cd1e2b59797154c7ec379ce3bb5d6779d9d78c +botocore==1.39.9 \ + --hash=sha256:02f141c2849e4589a79feea245ce4ecc478d48b7865572445af8aae3b041772d \ + --hash=sha256:a9691cbe03a3bc8b2720b3c36e5c5a2eecace6acd72bfb1107f00e75edaec4f3 # via # boto3 # libneuronxla # s3transfer -certifi==2025.4.26 \ - --hash=sha256:0a816057ea3cdefcef70270d2c515e4506bbc954f417fa5ade2021213bb8f0c6 \ - --hash=sha256:30350364dfe371162649852c63336a15c70c6510c2ad5015b21c2345311805f3 +certifi==2025.7.14 \ + --hash=sha256:6b31f564a415d79ee77df69d757bb49a5bb53bd9f756cbbe24394ffd6fc1f4b2 \ + --hash=sha256:8ea99dbdfaaf2ba2f9bac77b9249ef62ec5218e7c2b2e903378ed5fccf765995 # via requests charset-normalizer==3.4.2 \ --hash=sha256:005fa3432484527f9732ebd315da8da8001593e2cf46a3d817669f062c3d9ed4 \ @@ -169,8 +169,8 @@ jmespath==1.0.1 \ # via # boto3 # botocore -libneuronxla==2.2.3493.0+78c3e78c \ - --hash=sha256:9ba54e5d481afc25ae83e964a7c076543991ad80b88839cd500abcac130595e8 +libneuronxla==2.2.4410.0+835a67fb \ + --hash=sha256:37e5f08482ef1a9a844c2894de6d91c03400a092764e54ab67612456be439f16 # via -r runtimes/neuron/requirements.in lockfile==0.12.2 \ --hash=sha256:6aed02de03cba24efabcd600b30540140634fc06cfa603822d508d5361e9f799 \ @@ -206,8 +206,8 @@ networkx==2.8.8 \ --hash=sha256:230d388117af870fce5647a3c52401fcf753e94720e6ea6b4197a5355648885e \ --hash=sha256:e435dfa75b1d7195c7b8378c3859f0445cd88c6b0375c181ed66823a9ceb7524 # via neuronx-cc -neuronx-cc==2.18.121.0+9e31e41a \ - --hash=sha256:b4c0959594af716aff94698e15cc3997cc39a66c853c25a1bbf5eb7bf71fd7ec +neuronx-cc==2.19.8089.0+8ab9f450 \ + --hash=sha256:fbb05d8a0da8f54954c03ff2bea659f627309110c00efd54f523d877ddd6b88d # via # -r runtimes/neuron/requirements.in # libneuronxla @@ -287,20 +287,19 @@ python-dateutil==2.9.0.post0 \ --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \ --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427 # via botocore -requests==2.31.0 \ - --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \ - --hash=sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1 +requests==2.32.4 \ + --hash=sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c \ + --hash=sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422 # via # ec2-metadata - # neuronx-cc # requests-unixsocket requests-unixsocket==0.4.1 \ --hash=sha256:60c4942e9dbecc2f64d611039fb1dfc25da382083c6434ac0316dca3ff908f4d \ --hash=sha256:b2596158c356ecee68d27ba469a52211230ac6fb0cde8b66afb19f0ed47a1995 # via neuronx-cc -s3transfer==0.13.0 \ - --hash=sha256:0148ef34d6dd964d0d8cf4311b2b21c474693e57c2e069ec708ce043d2b527be \ - --hash=sha256:f5e6db74eb7776a37208001113ea7aa97695368242b364d73e91c981ac522177 +s3transfer==0.13.1 \ + --hash=sha256:a981aa7429be23fe6dfc13e80e4020057cbab622b08c0315288758d67cabc724 \ + --hash=sha256:c3fdba22ba1bd367922f27ec8032d6a1cf5f10c934fb5d68cf60fd5a23d936cf # via boto3 scipy==1.12.0 \ --hash=sha256:196ebad3a4882081f62a5bf4aeb7326aa34b110e533aab23e4374fcccb0890dc \ @@ -337,9 +336,13 @@ tqdm==4.67.1 \ --hash=sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2 \ --hash=sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2 # via neuronx-cc -urllib3==2.4.0 \ - --hash=sha256:414bc6535b787febd7567804cc015fee39daab8ad86268f1310a9250697de466 \ - --hash=sha256:4e16665048960a0900c702d4a66415956a584919c03361cac9f1df5c5dd7e813 +typing-extensions==4.14.1 \ + --hash=sha256:38b39f4aeeab64884ce9f74c94263ef78f3c22467c8724005483154c26648d36 \ + --hash=sha256:d1e1e3b58374dc93031d6eda2420a48ea44a36c2b4766a4fdeb3710755731d76 + # via neuronx-cc +urllib3==2.5.0 \ + --hash=sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760 \ + --hash=sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc # via # botocore # requests diff --git a/runtimes/neuron/zmlxneuron.c b/runtimes/neuron/zmlxneuron.c index ea5e11b..a5740be 100644 --- a/runtimes/neuron/zmlxneuron.c +++ b/runtimes/neuron/zmlxneuron.c @@ -6,7 +6,10 @@ void *zmlxneuron_dlopen(const char *filename, int flags) if (filename != NULL) { char *replacements[] = { - "libnccom.so", "libnccom.so.2", "libnccom-net.so", "libnccom-net.so.0", NULL, NULL, + "libnccom.so", "libnccom.so.2", + "libnrt.so", "libnrt.so.1", + "libncfw.so", "libncfw.so.2", + NULL, NULL, }; for (int i = 0; replacements[i] != NULL; i += 2) {