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

@@ -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")

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": {}

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,
+)

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,

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("@rules_zig//zig:defs.bzl", "BINARY_KIND", "zig_binary", "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")
 
@@ -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"],
-)
+    visibility = ["@libpjrt_neuron//:__subpackages__"],
-
-cc_shared_library(
-    name = "zmlxneuron_",
-    shared_lib_name = "libzmlxneuron.so.0",
-    deps = [":zmlxneuron_lib"],
-)
-
-cc_import(
-    name = "zmlxneuron",
-    shared_library = ":zmlxneuron_",
 )
 
 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"],
 )

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__"],
+)

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();
+}

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 = """\
+_BUILD_FILE_PRELUDE = """\
-load("@zml//bazel:cc_import.bzl", "cc_import")
+package(default_visibility = ["//visibility:public"])
 """
 
-_PACKAGES = {
+_UBUNTU_PACKAGES = {
-    "aws-neuronx-runtime-lib": packages.cc_import_glob_hdrs(
+    "zlib1g": packages.filegroup(name = "zlib1g", srcs = ["lib/x86_64-linux-gnu/libz.so.1"]),
-        name = "aws-neuronx-runtime-lib",
+    "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.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",
             },
-        visibility = ["@zml//runtimes/neuron:__subpackages__"],
-        deps = ["@aws-neuronx-collectives//:libnccom"],
         ),
+        packages.patchelf(
+            name = "libncfw.patchelf",
+            shared_library = "lib/libncfw.so",
+            soname = "libncfw.so.2",
+        ),
+    ]),
     "aws-neuronx-collectives": "\n".join([
-        packages.cc_import(
+        packages.filegroup(
             name = "libnccom",
-            shared_library = "lib/libnccom.so.2",
+            srcs = ["lib/libnccom.so.2"],
-            visibility = ["@aws-neuronx-runtime-lib//:__subpackages__"],
-        ),
-        packages.cc_import(
-            name = "libnccom-net",
-            shared_library = "lib/libnccom-net.so.0",
         ),
     ]),
 }
@@ -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(

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 {
+fn initialize(allocator: std.mem.Allocator, r_: *runfiles.Runfiles) !void {
-    var arena = std.heap.ArenaAllocator.init(std.heap.c_allocator);
-    const allocator = arena.allocator();
-    defer arena.deinit();
-
     {
         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();
+    try initialize(arena.allocator(), &r_);
-    return try asynk.callBlocking(pjrt.Api.loadFrom, .{"libpjrt_neuron.so"});
+
+    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});
+    };
 }

runtimes/neuron/packages.lock.json

@@ -143,22 +143,17 @@
 					"version": "12.2.0-14+deb12u1"
 				}
 			],
-			"key": "zlib1g_1-1.2.13.dfsg-1_amd64",
+			"key": "libgomp1_12.2.0-14-p-deb12u1_amd64",
-			"name": "zlib1g",
+			"name": "libgomp1",
-			"sha256": "d7dd1d1411fedf27f5e27650a6eff20ef294077b568f4c8c5e51466dc7c08ce4",
+			"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",
+			"key": "zlib1g_1-1.2.13.dfsg-1_amd64",
-			"name": "libxml2",
+			"name": "zlib1g",
-			"sha256": "35b76cb7038fc1c940204a4f05f33ffb79d027353ce469397d9adcf8f9b3e1a7",
+			"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",
+			"key": "aws-neuronx-runtime-lib_2.26.42.0-2ff3b5c7d_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",
 			"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",

runtimes/neuron/packages.yaml

@@ -18,8 +18,8 @@ archs:
     - "amd64"
 
 packages:
-    - "bash"
+    - "bash" # needed to launch neuronx-cc
-    - "zlib1g"
+    - "libgomp1" # implicit dep of neuronx-cc
-    - "libxml2"
+    - "zlib1g" # implicit dep of libpython
     - "aws-neuronx-runtime-lib"
     - "aws-neuronx-collectives"

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 \
+boto3==1.39.9 \
-    --hash=sha256:2f2cd517dd31d33ace0eefe567dc903fdf74221513e32f1e9445bdfac7554db7 \
+    --hash=sha256:5bc85e9fdec4e21ef5ca2c22b4d51a3e32b53f3da36ce51f5a3ea4dbde07b132 \
-    --hash=sha256:85c1556a110896f68de8573a9b4757c81071448dbf6ffc1074941bfc8a43195e
+    --hash=sha256:e3d3a6b617e1575e7ec854c820a882ab2e189a0421e74dc0dca2c9e13d4370a5
     # via libneuronxla
-botocore==1.38.25 \
+botocore==1.39.9 \
-    --hash=sha256:5a960bd990a11cdb78865e908a58ed712d87d9b419f55ad21c99d7d21f0d6726 \
+    --hash=sha256:02f141c2849e4589a79feea245ce4ecc478d48b7865572445af8aae3b041772d \
-    --hash=sha256:8c73e97d9662a6c92be33dab66cd1e2b59797154c7ec379ce3bb5d6779d9d78c
+    --hash=sha256:a9691cbe03a3bc8b2720b3c36e5c5a2eecace6acd72bfb1107f00e75edaec4f3
     # via
     #   boto3
     #   libneuronxla
     #   s3transfer
-certifi==2025.4.26 \
+certifi==2025.7.14 \
-    --hash=sha256:0a816057ea3cdefcef70270d2c515e4506bbc954f417fa5ade2021213bb8f0c6 \
+    --hash=sha256:6b31f564a415d79ee77df69d757bb49a5bb53bd9f756cbbe24394ffd6fc1f4b2 \
-    --hash=sha256:30350364dfe371162649852c63336a15c70c6510c2ad5015b21c2345311805f3
+    --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 \
+libneuronxla==2.2.4410.0+835a67fb \
-    --hash=sha256:9ba54e5d481afc25ae83e964a7c076543991ad80b88839cd500abcac130595e8
+    --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 \
+neuronx-cc==2.19.8089.0+8ab9f450 \
-    --hash=sha256:b4c0959594af716aff94698e15cc3997cc39a66c853c25a1bbf5eb7bf71fd7ec
+    --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 \
+requests==2.32.4 \
-    --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \
+    --hash=sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c \
-    --hash=sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1
+    --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 \
+s3transfer==0.13.1 \
-    --hash=sha256:0148ef34d6dd964d0d8cf4311b2b21c474693e57c2e069ec708ce043d2b527be \
+    --hash=sha256:a981aa7429be23fe6dfc13e80e4020057cbab622b08c0315288758d67cabc724 \
-    --hash=sha256:f5e6db74eb7776a37208001113ea7aa97695368242b364d73e91c981ac522177
+    --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 \
+typing-extensions==4.14.1 \
-    --hash=sha256:414bc6535b787febd7567804cc015fee39daab8ad86268f1310a9250697de466 \
+    --hash=sha256:38b39f4aeeab64884ce9f74c94263ef78f3c22467c8724005483154c26648d36 \
-    --hash=sha256:4e16665048960a0900c702d4a66415956a584919c03361cac9f1df5c5dd7e813
+    --hash=sha256:d1e1e3b58374dc93031d6eda2420a48ea44a36c2b4766a4fdeb3710755731d76
+    # via neuronx-cc
+urllib3==2.5.0 \
+    --hash=sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760 \
+    --hash=sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc
     # via
     #   botocore
     #   requests

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)
         {