Skip to content

Fix non-deterministic predictions caused by GC-invalidated binary addresses #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
pedroperrone wants to merge 3 commits into from
Closed

Fix non-deterministic predictions caused by GC-invalidated binary addresses #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0b0901c
Ensure tensor binary is in off-heap
pedroperrone Jul 16, 2026
7e9c65f
Keep references alive to prevent GC from clearing values while dirty …
pedroperrone Jul 16, 2026
e8d781d
Apply same fix to dmatrix module
pedroperrone Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 15 additions & 8 deletions lib/exgboost.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,8 @@ defmodule EXGBoost do

case data do
%Nx.Tensor{} = data ->
data_interface = ArrayInterface.from_tensor(data) |> Jason.encode!()
interface = ArrayInterface.from_tensor(data)
data_interface = Jason.encode!(interface)

{shape, preds} =
EXGBoost.NIF.booster_predict_from_dense(
Expand All @@ -244,30 +245,35 @@ defmodule EXGBoost do
)
|> Internal.unwrap!()

# Keep interface alive until after the NIF returns so its binary isn't
# freed while the NIF holds a raw pointer to it.
_keep_alive = interface.binary
Nx.tensor(preds) |> Nx.reshape(shape)

{%Nx.Tensor{} = indptr, %Nx.Tensor{} = indices, %Nx.Tensor{} = values, ncol} ->
indptr_interface = ArrayInterface.from_tensor(indptr) |> Jason.encode!()
indices_interface = ArrayInterface.from_tensor(indices) |> Jason.encode!()
values_interface = ArrayInterface.from_tensor(values) |> Jason.encode!()
indptr_interface = ArrayInterface.from_tensor(indptr)
indices_interface = ArrayInterface.from_tensor(indices)
values_interface = ArrayInterface.from_tensor(values)

{shape, preds} =
EXGBoost.NIF.booster_predict_from_csr(
boostr.ref,
indptr_interface,
indices_interface,
values_interface,
Jason.encode!(indptr_interface),
Jason.encode!(indices_interface),
Jason.encode!(values_interface),
ncol,
Jason.encode!(params),
proxy
)
|> Internal.unwrap!()

_keep_alive = {indptr_interface.binary, indices_interface.binary, values_interface.binary}
Nx.tensor(preds) |> Nx.reshape(shape)

data ->
data = Nx.concatenate(data)
data_interface = ArrayInterface.from_tensor(data) |> Jason.encode!()
interface = ArrayInterface.from_tensor(data)
data_interface = Jason.encode!(interface)

{shape, preds} =
EXGBoost.NIF.booster_predict_from_dense(
Expand All @@ -278,6 +284,7 @@ defmodule EXGBoost do
)
|> Internal.unwrap!()

_keep_alive = interface.binary
Nx.tensor(preds) |> Nx.reshape(shape)
end
end
Expand Down
11 changes: 10 additions & 1 deletion lib/exgboost/array_interface.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ defmodule EXGBoost.ArrayInterface do
"<#{Atom.to_string(tensor_type)}#{div(type_width, 8)}"
end

binary = Nx.to_binary(tensor)
binary = tensor |> Nx.to_binary() |> ensure_off_heap_binary()

tensor_addr =
EXGBoost.NIF.get_binary_address(binary) |> EXGBoost.Internal.unwrap!()
Expand All @@ -172,6 +172,15 @@ defmodule EXGBoost.ArrayInterface do
}
end

# Binaries ≤64 bytes are heap-allocated and can be moved by GC, invalidating
# any raw address captured by a NIF. Padding to >64 bytes forces a refc binary,
# which lives off-heap and has a stable address for the duration of the NIF call.
defp ensure_off_heap_binary(binary) when byte_size(binary) > 64, do: binary

defp ensure_off_heap_binary(binary) do
binary <> :binary.copy(<<0>>, 65 - byte_size(binary))
end

@spec get_tensor(EXGBoost.ArrayInterface.t()) :: Nx.Tensor.t()
def get_tensor(%__MODULE__{tensor: nil} = arr_int) do
num_items = arr_int.shape |> Tuple.to_list() |> Enum.product()
Expand Down
21 changes: 16 additions & 5 deletions lib/exgboost/dmatrix.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,13 +127,16 @@ defmodule EXGBoost.DMatrix do
args = Enum.into(Keyword.merge(meta_opts, str_opts), %{})

Enum.each(meta_opts, fn {key, value} ->
data_interface = ArrayInterface.from_tensor(value) |> Jason.encode!()
iface = ArrayInterface.from_tensor(value)
data_interface = Jason.encode!(iface)

EXGBoost.NIF.dmatrix_set_info_from_interface(
dmat.ref,
Atom.to_string(key),
data_interface
)

_keep_alive = iface.binary
end)

Enum.each(str_opts, fn {key, value} ->
Expand Down Expand Up @@ -311,13 +314,16 @@ defmodule EXGBoost.DMatrix do
config = Enum.into(config_opts, %{}, fn {key, value} -> {Atom.to_string(key), value} end)
format = Keyword.fetch!(format_opts, :format)

tensor_iface = ArrayInterface.from_tensor(tensor)

dmat =
EXGBoost.NIF.dmatrix_create_from_dense(
Jason.encode!(ArrayInterface.from_tensor(tensor)),
Jason.encode!(tensor_iface),
Jason.encode!(config)
)
|> Internal.unwrap!()

_keep_alive = tensor_iface.binary
set_params(%__MODULE__{ref: dmat, format: format}, opts)
end

Expand Down Expand Up @@ -370,17 +376,22 @@ defmodule EXGBoost.DMatrix do
raise ArgumentError, "Sparse format must be :csr or :csc"
end

indptr_iface = ArrayInterface.from_tensor(indptr)
indices_iface = ArrayInterface.from_tensor(indices)
data_iface = ArrayInterface.from_tensor(data)

dmat =
EXGBoost.NIF.dmatrix_create_from_sparse(
Jason.encode!(ArrayInterface.from_tensor(indptr)),
Jason.encode!(ArrayInterface.from_tensor(indices)),
Jason.encode!(ArrayInterface.from_tensor(data)),
Jason.encode!(indptr_iface),
Jason.encode!(indices_iface),
Jason.encode!(data_iface),
n,
Jason.encode!(config),
Atom.to_string(format)
)
|> Internal.unwrap!()

_keep_alive = {indptr_iface.binary, indices_iface.binary, data_iface.binary}
set_params(%__MODULE__{ref: dmat, format: format}, opts)
end
end
Expand Down
22 changes: 13 additions & 9 deletions test/nif_test.exs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,10 +90,14 @@ defmodule NifTest do
1.0
])

indptr_iface = from_tensor(indptr)
indices_iface = from_tensor(indices)
data_iface = from_tensor(data)

assert EXGBoost.NIF.dmatrix_create_from_sparse(
from_tensor(indptr) |> Jason.encode!(),
from_tensor(indices) |> Jason.encode!(),
from_tensor(data) |> Jason.encode!(),
Jason.encode!(indptr_iface),
Jason.encode!(indices_iface),
Jason.encode!(data_iface),
ncols,
config,
"csr"
Expand All @@ -102,9 +106,9 @@ defmodule NifTest do
:error

assert EXGBoost.NIF.dmatrix_create_from_sparse(
from_tensor(indptr) |> Jason.encode!(),
from_tensor(indices) |> Jason.encode!(),
from_tensor(data) |> Jason.encode!(),
Jason.encode!(indptr_iface),
Jason.encode!(indices_iface),
Jason.encode!(data_iface),
ncols,
config,
"csc"
Expand All @@ -114,9 +118,9 @@ defmodule NifTest do

{status, _} =
EXGBoost.NIF.dmatrix_create_from_sparse(
from_tensor(indptr) |> Jason.encode!(),
from_tensor(indices) |> Jason.encode!(),
from_tensor(data) |> Jason.encode!(),
Jason.encode!(indptr_iface),
Jason.encode!(indices_iface),
Jason.encode!(data_iface),
ncols,
config,
"csa"
Expand Down