diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 0000000..ed08691 --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,46 @@ +FROM hexpm/elixir:1.20.2-erlang-29.0.2-debian-trixie-20260623-slim + +ENV DEBIAN_FRONTEND=noninteractive +ENV LANG=C.UTF-8 +ENV LC_ALL=C.UTF-8 +ENV CC=gcc-14 +ENV CXX=g++-14 + +RUN apt-get update && \ + apt-get dist-upgrade --yes && \ + apt-get install --yes --no-install-recommends \ + build-essential \ + ca-certificates \ + cmake \ + curl \ + g++-14 \ + gcc-14 \ + git \ + inotify-tools \ + ninja-build \ + openssh-client \ + pkg-config \ + python3 \ + python3-venv \ + sudo \ + zsh && \ + update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 140 && \ + update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 140 && \ + update-alternatives --install /usr/bin/cc cc /usr/bin/gcc-14 140 && \ + update-alternatives --install /usr/bin/c++ c++ /usr/bin/g++-14 140 && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +RUN useradd -m -s /bin/zsh -u 1000 -G sudo dev && \ + echo 'dev ALL=(ALL) NOPASSWD:ALL' > /etc/sudoers.d/dev && \ + chmod 0440 /etc/sudoers.d/dev && \ + mkdir -p /workspace && \ + chown dev:dev /workspace + +USER dev + +RUN git config --global --add safe.directory /workspace && \ + mix local.hex --force && \ + mix local.rebar --force + +WORKDIR /workspace diff --git a/.devcontainer/devcontainer-lock.json b/.devcontainer/devcontainer-lock.json new file mode 100644 index 0000000..523de18 --- /dev/null +++ b/.devcontainer/devcontainer-lock.json @@ -0,0 +1,24 @@ +{ + "features": { + "ghcr.io/devcontainers/features/common-utils:2": { + "version": "2.5.9", + "resolved": "ghcr.io/devcontainers/features/common-utils@sha256:cb0c4d3c276f157eed17935747e364178d75fee17f55c4e129966f64633deb3a", + "integrity": "sha256:cb0c4d3c276f157eed17935747e364178d75fee17f55c4e129966f64633deb3a" + }, + "ghcr.io/devcontainers/features/docker-outside-of-docker:1": { + "version": "1.10.0", + "resolved": "ghcr.io/devcontainers/features/docker-outside-of-docker@sha256:c2c2cf829505ead8e4892c88c31b6594ae94a2bbb209e16e1fac456c1a3a624e", + "integrity": "sha256:c2c2cf829505ead8e4892c88c31b6594ae94a2bbb209e16e1fac456c1a3a624e" + }, + "ghcr.io/devcontainers/features/git:1": { + "version": "1.3.7", + "resolved": "ghcr.io/devcontainers/features/git@sha256:a3e43ff91b9f5f6bd2c14bd510d43e5e698f1266dc41027ba4e04e7e45be607a", + "integrity": "sha256:a3e43ff91b9f5f6bd2c14bd510d43e5e698f1266dc41027ba4e04e7e45be607a" + }, + "ghcr.io/devcontainers/features/github-cli:1": { + "version": "1.1.0", + "resolved": "ghcr.io/devcontainers/features/github-cli@sha256:d22f50b70ed75339b4eed1ba9ecde3a1791f90e88d37936517e3bace0bbad671", + "integrity": "sha256:d22f50b70ed75339b4eed1ba9ecde3a1791f90e88d37936517e3bace0bbad671" + } + } +} diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 0000000..c452d66 --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,33 @@ +{ + "name": "EXGBoost Erlang 29 GCC 14", + "build": { + "dockerfile": "Dockerfile", + "context": ".." + }, + "workspaceMount": "source=${localWorkspaceFolder},target=/workspace,type=bind,consistency=cached", + "workspaceFolder": "/workspace", + "containerEnv": { + "CC": "gcc-14", + "CXX": "g++-14", + "MIX_ENV": "dev" + }, + "customizations": { + "vscode": { + "extensions": [ + "jakebecker.elixir-ls" + ] + } + }, + "postCreateCommand": "mix deps.get", + "remoteUser": "dev", + "features": { + "ghcr.io/devcontainers/features/common-utils:2": { + "configureZshAsDefaultShell": true + }, + "ghcr.io/devcontainers/features/git:1": {}, + "ghcr.io/devcontainers/features/github-cli:1": {}, + "ghcr.io/devcontainers/features/docker-outside-of-docker:1": { + "moby": false + } + } +} diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..5ae3052 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,12 @@ +version: 2 +updates: + - package-ecosystem: "mix" + directory: "/" + schedule: + interval: "weekly" + open-pull-requests-limit: 10 + + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" \ No newline at end of file diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..1c0bf8d --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,53 @@ +name: ci + +on: + pull_request: + push: + branches: + - main + +jobs: + test: + name: Test OTP ${{ matrix.otp }} / Elixir ${{ matrix.elixir }} + runs-on: ubuntu-24.04 + env: + CC: gcc-14 + CXX: g++-14 + MIX_ENV: test + strategy: + fail-fast: false + matrix: + include: + - otp: "28.3" + elixir: "1.18.4" + - otp: "28.3" + elixir: "1.19.5" + - otp: "29.0.2" + elixir: "1.20.2" + + steps: + - uses: actions/checkout@v5 + + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install -y build-essential cmake ninja-build gcc-14 g++-14 + + - uses: erlef/setup-beam@v1 + with: + otp-version: ${{ matrix.otp }} + elixir-version: ${{ matrix.elixir }} + + - name: Install Hex and Rebar + run: | + mix local.hex --force + mix local.rebar --force + + - name: Install dependencies + run: mix deps.get + + - name: Check formatting + run: mix format --check-formatted + + - name: Run tests + run: MIX_ENV=test mix test diff --git a/.github/workflows/precompile.yml b/.github/workflows/precompile.yml index dfbae0c..1494c92 100644 --- a/.github/workflows/precompile.yml +++ b/.github/workflows/precompile.yml @@ -10,20 +10,32 @@ jobs: runs-on: ubuntu-24.04 env: MIX_ENV: "prod" + CC: gcc-14 + CXX: g++-14 strategy: matrix: - otp: ["27.0", "28.0"] - elixir: ["1.18.4"] + include: + - otp: "28.3" + elixir: "1.18.4" + - otp: "28.3" + elixir: "1.19.5" + - otp: "29.0.2" + elixir: "1.20.2" steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v5 - uses: erlef/setup-beam@v1 with: otp-version: ${{matrix.otp}} elixir-version: ${{matrix.elixir}} - - name: Install system dependecies + - name: Install Hex and Rebar + run: | + mix local.hex --force + mix local.rebar --force + - name: Install system dependencies run: | sudo apt-get update sudo apt-get install -y build-essential automake autoconf pkg-config bc m4 unzip zip \ + cmake ninja-build gcc-14 g++-14 \ gcc-aarch64-linux-gnu g++-aarch64-linux-gnu \ gcc-riscv64-linux-gnu g++-riscv64-linux-gnu - name: Mix Test @@ -42,25 +54,16 @@ jobs: cache/*.tar.gz macos: - runs-on: ${{matrix.runner}} - # Homebrew supports versioned Erlang/OTP but not Elixir - # It's a deliberate design decision from Homebrew to - # only support versioned distributions for certain packages - name: Mac (ARM) Erlang/OTP ${{matrix.otp}} / Elixir + runs-on: macos-15 + name: Mac (ARM) Erlang/OTP 29.0.2 / Elixir 1.20.2 env: MIX_ENV: "prod" - strategy: - matrix: - runner: ["macos-15"] - otp: ["27.0", "28.0"] - elixir: ["1.18.4"] steps: - - uses: actions/checkout@v3 - - uses: asdf-vm/actions/install@v2 + - uses: actions/checkout@v5 + - uses: erlef/setup-beam@v1 with: - tool_versions: | - erlang ${{matrix.otp}} - elixir ${{matrix.elixir}} + otp-version: "29.0.2" + elixir-version: "1.20.2" - name: Install libomp run: | brew install libomp diff --git a/.gitignore b/.gitignore index 15fd0cd..839a055 100644 --- a/.gitignore +++ b/.gitignore @@ -9,7 +9,6 @@ erl_crash.dump *.beam /config/*.secret.exs .elixir_ls/ -.tool-versions .vscode/ checksum.exs .DS_Store \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..d11fe45 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,68 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## 0.10.0 + +### Added + +- Comprehensive safety validation tests in `test/safety_improvements_test.exs` covering: + - Invalid typestr format rejection + - Binary size validation preventing buffer overflows + - Boolean type strictness (only `true`/`false` atoms accepted) + - Shape overflow protection for extremely large dimensions + - Endianness marker validation + +### Changed + +- **NIF Layer Safety Improvements**: + - Buffer overflow protection: Dynamic allocation with overflow-checked arithmetic replaces fixed-size buffers in Array Interface JSON builder + - Strict type safety: All pointer arithmetic now uses `uintptr_t` with `PRIuPTR` formatting + - Platform-independent integer conversions: `enif_get_ulong`/`enif_make_ulong` replaced with `enif_get_uint64`/`enif_make_uint64` + - Comprehensive input validation: Added validation for typestr format, shape dimensions, boolean values, and binary sizes + - Memory safety: Binary size validation ensures data buffers match expected sizes before any memory operations + - Single cleanup path: Refactored to use consistent resource management with single cleanup labels to prevent memory leaks + - Atomic data copying: All data returned from XGBoost is now copied atomically within NIF calls, eliminating pointer lifetime issues + - Added static assertions and overflow checks for VLA (Variable Length Array) declarations + +- **ArrayInterface Optimization**: + - Removed `address` field - pointer addresses are never exposed to Elixir for safety + - Removed `tensor` field - eliminates duplicate memory storage; tensors now reconstructed on-demand from binary data + - Removed `Jason.Encoder` protocol implementation - no longer needed as ArrayInterface is not serialized to JSON + - Updated `Inspect` protocol to show `readonly` directly instead of `data: [address, readonly]` + - `get_tensor/1` now reconstructs tensors from binary data instead of caching, trading minimal performance for significant memory savings + - Updated `from_map/1` to ignore address values from incoming data, only extracting readonly flag + +- **DMatrix Improvements**: + - `get_quantile_cut/1` refactored to return maps with `:binary`, `:typestr`, `:shape` instead of JSON with memory addresses + - Added internal `build_tensor_from_map/1` helper to reconstruct tensors from NIF-returned data + - All data copying happens atomically within C before returning to Elixir + +- **Unsafe APIs marked as DEPRECATED - to be removed in future releases**: + - `EXGBoost.NIF.get_binary_from_address/2` - arbitrary memory read primitive that could crash the BEAM VM + - `exg_get_binary_from_address` C NIF function and all declarations + - Address-based tensor reconstruction in `ArrayInterface.get_tensor/1` + - Tensor caching in `ArrayInterface` struct + +### Fixed + +- Cross-platform compatibility: Fixed `unsigned long` → `ErlNifUInt64` conversions to prevent data truncation on Windows (LP64 vs LLP64 calling conventions) +- Memory safety: Eliminated all pointer lifetime gaps where Elixir code held addresses to freed memory +- Buffer safety: All shape-to-JSON conversions now use dynamic allocation with proper bounds checking + +### Security + +- Eliminated arbitrary memory read primitive that allowed reading from any address +- All binary data is now validated for size before access, preventing buffer overflows +- Pointer addresses are never exposed to Elixir, preventing use-after-free vulnerabilities +- Strict type validation prevents type confusion attacks + +## [0.9.1] + +### Removed + +- Kino and Livebook Integration; `kino` doesn't seem like it's under active development, make this a pure library. +- Mark `EXGBoost.DMatrix.from_file` as deprecated. diff --git a/Makefile b/Makefile index 1768938..2fd89e0 100644 --- a/Makefile +++ b/Makefile @@ -1,17 +1,28 @@ # Environment variables passed via elixir_make # ERTS_INCLUDE_DIR # MIX_APP_PATH +# MIX_ENV TEMP ?= $(HOME)/.cache +MIX_ENV ?= dev XGBOOST_CACHE ?= $(TEMP)/exgboost XGBOOST_GIT_REPO ?= https://github.com/dmlc/xgboost.git -# v3.0.5 tagged release -XGBOOST_GIT_REV ?= v3.0.5 +# v3.1.3 tagged release +XGBOOST_GIT_REV ?= v3.1.3 +OLD_XGBOOST_GIT_REV ?= v3.0.5 +NEW_XGBOOST_GIT_REV ?= $(XGBOOST_GIT_REV) XGBOOST_NS = xgboost-$(XGBOOST_GIT_REV) XGBOOST_DIR = $(XGBOOST_CACHE)/$(XGBOOST_NS) XGBOOST_LIB_DIR = $(XGBOOST_DIR)/build/xgboost XGBOOST_LIB_DIR_FLAG = $(XGBOOST_LIB_DIR)/exgboost.ok +# Set build type based on MIX_ENV +ifeq ($(MIX_ENV), prod) + CMAKE_BUILD_TYPE = Release +else + CMAKE_BUILD_TYPE = RelWithDebInfo +endif + # Private configuration PRIV_DIR = $(MIX_APP_PATH)/priv EXGBOOST_DIR = $(realpath c/exgboost) @@ -43,21 +54,21 @@ else endif $(EXGBOOST_SO): $(EXGBOOST_CACHE_SO) - @ mkdir -p $(PRIV_DIR) + @mkdir -p $(PRIV_DIR) cp -a $(abspath $(EXGBOOST_CACHE_LIB_DIR)) $(EXGBOOST_LIB_DIR) ; \ cp -a $(abspath $(EXGBOOST_CACHE_SO)) $(EXGBOOST_SO) ; $(EXGBOOST_CACHE_SO): $(XGBOOST_LIB_DIR_FLAG) $(C_SRCS) @mkdir -p cache - cp -a $(XGBOOST_LIB_DIR) $(EXGBOOST_CACHE_LIB_DIR) - mv $(XGBOOST_LIB_DIR)/lib/$(LIBXGBOOST) $(EXGBOOST_CACHE_LIB_DIR) + cp -R $(XGBOOST_LIB_DIR) $(EXGBOOST_CACHE_LIB_DIR) + cp $(XGBOOST_DIR)/lib/$(LIBXGBOOST) $(EXGBOOST_CACHE_LIB_DIR) $(CC) $(CFLAGS) $(wildcard $(EXGBOOST_DIR)/src/*.c) $(LDFLAGS) -o $(EXGBOOST_CACHE_SO) $(POST_INSTALL) # This new target handles fetching the source code. # It only runs if the .git directory inside the source folder is missing. $(XGBOOST_DIR)/.git: - mkdir -p $(XGBOOST_DIR) && \ + @mkdir -p $(XGBOOST_DIR) && \ cd $(XGBOOST_DIR) && \ git init && \ git remote add origin $(XGBOOST_GIT_REPO) && \ @@ -69,10 +80,30 @@ $(XGBOOST_DIR)/.git: # It only contains the build commands. $(XGBOOST_LIB_DIR_FLAG): $(XGBOOST_DIR)/.git cd $(XGBOOST_DIR) && \ - cmake -B build -S . -DCMAKE_INSTALL_PREFIX=$(XGBOOST_LIB_DIR) -DCMAKE_BUILD_TYPE=RelWithDebInfo -GNinja $(CMAKE_FLAGS) && \ + cmake -B build -S . -DCMAKE_INSTALL_PREFIX=$(XGBOOST_LIB_DIR) -DCMAKE_BUILD_TYPE=$(CMAKE_BUILD_TYPE) -GNinja $(CMAKE_FLAGS) && \ ninja -C build install touch $(XGBOOST_LIB_DIR_FLAG) +check-xgboost-c-api: $(XGBOOST_LIB_DIR_FLAG) + ./scripts/check_xgboost_c_api.sh "$(XGBOOST_LIB_DIR)/include" + +compare-xgboost-c-api: + @set -eu; \ + for rev in "$(OLD_XGBOOST_GIT_REV)" "$(NEW_XGBOOST_GIT_REV)"; do \ + dir="$(XGBOOST_CACHE)/xgboost-$$rev"; \ + mkdir -p "$$dir"; \ + if [ ! -d "$$dir/.git" ]; then \ + git -C "$$dir" init; \ + git -C "$$dir" remote add origin "$(XGBOOST_GIT_REPO)"; \ + fi; \ + git -C "$$dir" fetch --depth 1 --recurse-submodules origin "$$rev"; \ + git -C "$$dir" checkout -f FETCH_HEAD; \ + git -C "$$dir" submodule update --init --recursive; \ + done; \ + ./scripts/check_xgboost_c_api.sh --compare \ + "$(XGBOOST_CACHE)/xgboost-$(OLD_XGBOOST_GIT_REV)/include" \ + "$(XGBOOST_CACHE)/xgboost-$(NEW_XGBOOST_GIT_REV)/include" + clean: rm -rf $(EXGBOOST_CACHE_SO) rm -rf $(EXGBOOST_CACHE_LIB_DIR) diff --git a/README.md b/README.md index e11c5f1..3f810f3 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,7 @@ EXGBoost.predict(model, x) EXGBoost is designed to feel familiar to the users of the Python XGBoost library. `EXGBoost.train/2` is the primary entry point for training a model. It accepts a Nx tensor for the features and a Nx tensor for the labels. -`EXGBoost.train/2` returns a trained`Booster` struct that can be used for prediction. `EXGBoost.train/2` also +`EXGBoost.train/2` returns a trained `Booster` struct that can be used for prediction. `EXGBoost.train/2` also accepts a keyword list of options that can be used to configure the training process. See the [XGBoost documentation](https://xgboost.readthedocs.io/en/latest/parameter.html) for the full list of options. @@ -119,6 +119,17 @@ It accepts a `Booster` struct (which is the output of `EXGBoost.train/2`). preds = EXGBoost.train(X, y) |> EXGBoost.predict(X) ``` +## Concurrency and Thread Safety + +**Important**: Booster objects are **not thread-safe** for concurrent predictions. The underlying XGBoost C API does not provide synchronization mechanisms, and sharing a single booster reference across multiple Elixir processes for concurrent predictions can lead to race conditions, memory corruption, or incorrect results. For this reason it is not recommended that you cache boosters to be used +by multiple tasks in calling applications. + +### Why This Matters + +- `EXGBoost.predict/2` and `EXGBoost.inplace_predict/2` both use dirty CPU-bound NIF schedulers +- This prevents blocking the BEAM scheduler but **does not** provide thread safety +- Concurrent access to the same booster from multiple processes can cause undefined behavior + ## Serialization A Booster can be serialized to a file using `EXGBoost.write_*` and loaded from a file @@ -155,17 +166,6 @@ preds = EXGBoost.train(X, y) |> EXGBoost.predict(X) You can see available styles by running `EXGBoost.Plotting.get_styles()` or refer to the `EXGBoost.Plotting.Styles` documentation for a gallery of the styles. -## Kino & Livebook Integration - - `EXGBoost` integrates with [Kino](https://hexdocs.pm/kino/Kino.html) and [Livebook](https://livebook.dev/) - to provide a rich interactive experience for data scientists. - - EXGBoost implements the `Kino.Render` protocol for `EXGBoost.Booster` structs. This allows you to render - a Booster in a Livebook notebook. Under the hood, `EXGBoost` uses [Vega-Lite](https://vega.github.io/vega-lite/) - and [Kino Vega-Lite](https://hexdocs.pm/kino_vega_lite/Kino.VegaLite.html) to render the Booster. - - See the [`Plotting in EXGBoost`](notebooks/plotting.livemd) Notebook for an example of how to use `EXGBoost` with `Kino` and `Livebook`. - ## Examples See the example Notebooks in the left sidebar (under the `Pages` tab) for more examples and tutorials @@ -175,7 +175,7 @@ preds = EXGBoost.train(X, y) |> EXGBoost.predict(X) ### Precompiled Distribution -We currenly offer the following precompiled packages for EXGBoost: +We currently offer the following precompiled packages for EXGBoost: ```elixir %{ @@ -220,3 +220,21 @@ You also need to set `CC_PRECOMPILER_PRECOMPILE_ONLY_LOCAL=true` before the firs ## License Licensed under an [Apache-2](https://github.com/acalejos/exgboost/blob/main/LICENSE) license. + +## Maintaining this fork + +Recommended fork maintenance cadence: + +- Keep dependencies current with Dependabot PRs (configured in `.github/dependabot.yml`). +- Ensure all PRs pass CI (`.github/workflows/ci.yml`) before merge. +- Keep precompiled artefacts and checksums aligned with each release tag (`.github/workflows/precompile.yml`). +- When upgrading `nx`, verify supported Elixir/OTP versions in CI matrix and `mix.exs`. + +Native compatibility checks: + +- Run `make check-xgboost-c-api` after changing `XGBOOST_GIT_REV` or native C files. +- Run `mix test test/nif_test.exs` to validate runtime behavior after the API check passes. +- To compare API declarations between two XGBoost versions, run: + `scripts/check_xgboost_c_api.sh --compare ` +- Or use make to fetch and compare tags directly: + `make compare-xgboost-c-api OLD_XGBOOST_GIT_REV=v3.0.5 NEW_XGBOOST_GIT_REV=v3.1.3` diff --git a/c/exgboost/include/utils.h b/c/exgboost/include/utils.h index a995657..92f987d 100644 --- a/c/exgboost/include/utils.h +++ b/c/exgboost/include/utils.h @@ -2,6 +2,8 @@ #define EXGBOOST_UTILS_H #include +#include +#include #include #include #include @@ -10,6 +12,9 @@ ErlNifResourceType *DMatrix_RESOURCE_TYPE; ErlNifResourceType *Booster_RESOURCE_TYPE; typedef uint64_t bst_ulong; +// Initialize atoms (must be called during NIF load) +void exg_init_atoms(ErlNifEnv *env); + void DMatrix_RESOURCE_TYPE_cleanup(ErlNifEnv *env, void *arg); void Booster_RESOURCE_TYPE_cleanup(ErlNifEnv *env, void *arg); @@ -22,9 +27,11 @@ ERL_NIF_TERM ok_atom(ErlNifEnv *env); ERL_NIF_TERM exg_ok(ErlNifEnv *env, ERL_NIF_TERM term); +// Unsafe, to be deprecated in future releases. ERL_NIF_TERM exg_get_binary_address(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]); +// Unsafe, to be deprecated in future releases. ERL_NIF_TERM exg_get_binary_from_address(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]); @@ -42,4 +49,14 @@ int exg_get_string_list(ErlNifEnv *env, ERL_NIF_TERM term, char ***out, int exg_get_dmatrix_list(ErlNifEnv *env, ERL_NIF_TERM term, DMatrixHandle **dmats, unsigned *len); +void exg_free_string_list(char **items, unsigned len); + +void exg_free_dmatrix_list(DMatrixHandle *dmats); + +// Array Interface helper - builds JSON from components with fresh address +int exg_build_array_interface_json(ErlNifEnv *env, ERL_NIF_TERM binary_term, + ERL_NIF_TERM typestr_term, ERL_NIF_TERM shape_term, + ERL_NIF_TERM readonly_term, char **json_out, + const char **error_msg); + #endif \ No newline at end of file diff --git a/c/exgboost/src/booster.c b/c/exgboost/src/booster.c index 7f578f1..64f707e 100644 --- a/c/exgboost/src/booster.c +++ b/c/exgboost/src/booster.c @@ -7,6 +7,7 @@ static ERL_NIF_TERM make_Booster_resource(ErlNifEnv *env, enif_alloc_resource(Booster_RESOURCE_TYPE, sizeof(BoosterHandle *)); if (resource != NULL) { *resource = handle; + // BEAM resource now owns the handle and releases it in resource cleanup. ret = exg_ok(env, enif_make_resource(env, resource)); enif_release_resource(resource); } else { @@ -49,6 +50,8 @@ ERL_NIF_TERM EXGBoosterCreate(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); } END: + // exg_get_dmatrix_list allocates this temporary array. + exg_free_dmatrix_list(dmats); return ret; } @@ -148,6 +151,7 @@ ERL_NIF_TERM EXGBoosterSetParam(ErlNifEnv *env, int argc, ret = exg_error(env, "Booster parameter value must be a string"); goto END; } + // XGBoost consumes name/value during this call; no ownership transfer. result = XGBoosterSetParam(booster, name, value); if (result == 0) { ret = enif_make_atom(env, "ok"); @@ -155,6 +159,12 @@ ERL_NIF_TERM EXGBoosterSetParam(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); } END: + if (name != NULL) { + enif_free(name); + } + if (value != NULL) { + enif_free(value); + } return ret; } @@ -177,7 +187,7 @@ ERL_NIF_TERM EXGBoosterGetNumFeature(ErlNifEnv *env, int argc, booster = *resource; result = XGBoosterGetNumFeature(booster, &num_feature); if (result == 0) { - ret = exg_ok(env, enif_make_ulong(env, num_feature)); + ret = exg_ok(env, enif_make_uint64(env, (ErlNifUInt64)num_feature)); } else { ret = exg_error(env, XGBGetLastError()); } @@ -290,7 +300,7 @@ ERL_NIF_TERM EXGBoosterEvalOneIter(ErlNifEnv *env, int argc, int iter = -1; unsigned num_dmats = 0; unsigned num_evnames = 0; - char *out = NULL; + const char *out = NULL; ERL_NIF_TERM ret = -1; int result = -1; if (4 != argc) { @@ -319,7 +329,7 @@ ERL_NIF_TERM EXGBoosterEvalOneIter(ErlNifEnv *env, int argc, ret = exg_error(env, "dmats and evnames must have the same length"); goto END; } - result = XGBoosterEvalOneIter(booster, iter, dmats, evnames, + result = XGBoosterEvalOneIter(booster, iter, dmats, (const char **)evnames, (bst_ulong)num_dmats, &out); if (result == 0) { ret = exg_ok(env, enif_make_string(env, out, ERL_NIF_LATIN1)); @@ -327,6 +337,9 @@ ERL_NIF_TERM EXGBoosterEvalOneIter(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); } END: + // Helper-allocated arrays must be reclaimed on all paths. + exg_free_dmatrix_list(dmats); + exg_free_string_list(evnames, num_evnames); return ret; } @@ -335,7 +348,7 @@ ERL_NIF_TERM EXGBoosterGetAttr(ErlNifEnv *env, int argc, BoosterHandle booster; BoosterHandle **booster_resource = NULL; char *key = NULL; - char *out = NULL; + const char *out = NULL; ERL_NIF_TERM ret = -1; int result = -1; int success = -1; @@ -364,6 +377,9 @@ ERL_NIF_TERM EXGBoosterGetAttr(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); } END: + if (key != NULL) { + enif_free(key); + } return ret; } @@ -415,6 +431,12 @@ ERL_NIF_TERM EXGBoosterSetAttr(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); } END: + if (key != NULL) { + enif_free(key); + } + if (value != NULL) { + enif_free(value); + } return ret; } @@ -422,7 +444,7 @@ ERL_NIF_TERM EXGBoosterGetAttrNames(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { BoosterHandle booster; BoosterHandle **booster_resource = NULL; - char **out = NULL; + const char **out = NULL; bst_ulong out_len = 0; ERL_NIF_TERM ret = -1; int result = -1; @@ -438,14 +460,16 @@ ERL_NIF_TERM EXGBoosterGetAttrNames(ErlNifEnv *env, int argc, booster = *booster_resource; result = XGBoosterGetAttrNames(booster, &out_len, &out); if (result == 0) { + // Check VLA size fits in size_t + if (out_len > SIZE_MAX / sizeof(ERL_NIF_TERM)) { + ret = exg_error(env, "Result is too large"); + goto END; + } ERL_NIF_TERM arr[out_len]; for (bst_ulong i = 0; i < out_len; ++i) { - char *local = enif_alloc(strlen(out[i]) + 1); - strcpy(local, out[i]); - arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1); - // TODO: Do we free here or is it handled by the XGBoost library / BEAM? + arr[i] = enif_make_string(env, out[i], ERL_NIF_LATIN1); } - ret = exg_ok(env, enif_make_list_from_array(env, arr, out_len)); + ret = exg_ok(env, enif_make_list_from_array(env, arr, (size_t)out_len)); } else { ret = exg_error(env, XGBGetLastError()); } @@ -485,7 +509,8 @@ ERL_NIF_TERM EXGBoosterSetStrFeatureInfo(ErlNifEnv *env, int argc, goto END; } handle = *resource; - result = XGBoosterSetStrFeatureInfo(handle, field, features, num_features); + result = XGBoosterSetStrFeatureInfo(handle, field, (const char **)features, + num_features); if (result == 0) { ret = ok_atom(env); } else { @@ -493,8 +518,10 @@ ERL_NIF_TERM EXGBoosterSetStrFeatureInfo(ErlNifEnv *env, int argc, } END: if (features != NULL) { - enif_free(features); - features = NULL; + exg_free_string_list(features, num_features); + } + if (field != NULL) { + enif_free(field); } return ret; } @@ -529,14 +556,17 @@ ERL_NIF_TERM EXGBoosterGetStrFeatureInfo(ErlNifEnv *env, int argc, result = XGBoosterGetStrFeatureInfo(handle, field, &out_size, &c_out_features); if (result == 0) { + // Check VLA size fits in size_t + if (out_size > SIZE_MAX / sizeof(ERL_NIF_TERM)) { + ret = exg_error(env, "Result is too large"); + goto END; + } ERL_NIF_TERM arr[out_size]; for (bst_ulong i = 0; i < out_size; ++i) { - char *local = enif_alloc(strlen(c_out_features[i]) + 1); - strcpy(local, c_out_features[i]); - arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1); - // TODO: Do we free here or is it handled by the XGBoost library / BEAM? + // enif_make_string materializes a BEAM term; no temporary C copy needed. + arr[i] = enif_make_string(env, c_out_features[i], ERL_NIF_LATIN1); } - ret = exg_ok(env, enif_make_list_from_array(env, arr, out_size)); + ret = exg_ok(env, enif_make_list_from_array(env, arr, (size_t)out_size)); } else { ret = exg_error(env, XGBGetLastError()); } @@ -551,12 +581,12 @@ ERL_NIF_TERM EXGBoosterFeatureScore(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { BoosterHandle booster; BoosterHandle **booster_resource = NULL; - char **config = NULL; + char *config = NULL; bst_ulong out_n_features = 0; - char **out_features = NULL; + const char **out_features = NULL; bst_ulong out_dim = 0; - bst_ulong *out_shape = NULL; - float *out_scores = NULL; + const bst_ulong *out_shape = NULL; + const float *out_scores = NULL; ERL_NIF_TERM ret = -1; int result = -1; if (2 != argc) { @@ -577,6 +607,12 @@ ERL_NIF_TERM EXGBoosterFeatureScore(ErlNifEnv *env, int argc, XGBoosterFeatureScore(booster, config, &out_n_features, &out_features, &out_dim, &out_shape, &out_scores); if (result == 0) { + // Check VLA sizes fit in size_t + if (out_n_features > SIZE_MAX / sizeof(ERL_NIF_TERM) || + out_dim > SIZE_MAX / sizeof(ERL_NIF_TERM)) { + ret = exg_error(env, "Result is too large"); + goto END; + } ERL_NIF_TERM feature_arr[out_n_features]; for (bst_ulong i = 0; i < out_n_features; ++i) { ERL_NIF_TERM shape_arr[out_dim]; @@ -596,13 +632,16 @@ ERL_NIF_TERM EXGBoosterFeatureScore(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); } END: + if (config != NULL) { + enif_free(config); + } return ret; } static ERL_NIF_TERM collect_prediction_results(ErlNifEnv *env, - bst_ulong *out_shape, + const bst_ulong *out_shape, bst_ulong out_dim, - float *out_result) { + const float *out_result) { bst_ulong out_len = 1; ERL_NIF_TERM shape_arr[out_dim]; for (bst_ulong j = 0; j < out_dim; ++j) { @@ -612,6 +651,7 @@ static ERL_NIF_TERM collect_prediction_results(ErlNifEnv *env, ERL_NIF_TERM shape = enif_make_tuple_from_array(env, shape_arr, out_dim); ERL_NIF_TERM result_arr[out_len]; for (bst_ulong i = 0; i < out_len; ++i) { + // Values are copied into BEAM-managed terms here. result_arr[i] = enif_make_double(env, out_result[i]); } return exg_ok(env, enif_make_tuple2( @@ -626,9 +666,9 @@ ERL_NIF_TERM EXGBoosterPredictFromDMatrix(ErlNifEnv *env, int argc, DMatrixHandle dmatrix; DMatrixHandle **dmatrix_resource = NULL; char *config = NULL; - bst_ulong *out_shape = NULL; + const bst_ulong *out_shape = NULL; bst_ulong out_dim = 0; - float *out_result = NULL; + const float *out_result = NULL; ERL_NIF_TERM ret = -1; int result = -1; @@ -673,30 +713,37 @@ ERL_NIF_TERM EXGBoosterPredictFromDense(ErlNifEnv *env, int argc, DMatrixHandle proxy; DMatrixHandle **proxy_resource = NULL; char *values = NULL; + const char *error_msg = NULL; char *config = NULL; - bst_ulong *out_shape = NULL; + const bst_ulong *out_shape = NULL; bst_ulong out_dim = 0; - float *out_result = NULL; + const float *out_result = NULL; int result = -1; ERL_NIF_TERM ret = -1; - if (4 != argc) { + + if (argc != 7) { ret = exg_error(env, "Wrong number of arguments"); goto END; } + if (!enif_get_resource(env, argv[0], Booster_RESOURCE_TYPE, (void *)&(booster_resource))) { ret = exg_error(env, "Invalid Booster"); goto END; } - if (!exg_get_string(env, argv[1], &values)) { - ret = exg_error(env, "Value must be a JSON-encoded string"); + + // Build ArrayInterface JSON from components: (binary, typestr, shape, readonly) + if (!exg_build_array_interface_json(env, argv[1], argv[2], argv[3], argv[4], &values, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract ArrayInterface"); goto END; } - if (!exg_get_string(env, argv[2], &config)) { + + if (!exg_get_string(env, argv[5], &config)) { ret = exg_error(env, "Config must be a JSON-encoded string"); goto END; } - if (!enif_get_resource(env, argv[3], DMatrix_RESOURCE_TYPE, + + if (!enif_get_resource(env, argv[6], DMatrix_RESOURCE_TYPE, (void *)&(proxy_resource))) { proxy = NULL; } else { @@ -719,6 +766,7 @@ ERL_NIF_TERM EXGBoosterPredictFromDense(ErlNifEnv *env, int argc, } return ret; } + ERL_NIF_TERM EXGBoosterPredictFromCSR(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { BoosterHandle booster; @@ -728,43 +776,58 @@ ERL_NIF_TERM EXGBoosterPredictFromCSR(ErlNifEnv *env, int argc, char *indptr = NULL; char *indices = NULL; char *data = NULL; + const char *error_msg = NULL; char *config = NULL; - bst_ulong ncols = 0; - bst_ulong *out_shape = NULL; + int ncols = 0; + const bst_ulong *out_shape = NULL; bst_ulong out_dim = 0; - float *out_result = NULL; + const float *out_result = NULL; int result = -1; ERL_NIF_TERM ret = -1; - if (7 != argc) { + + if (argc != 16) { ret = exg_error(env, "Wrong number of arguments"); goto END; } + if (!enif_get_resource(env, argv[0], Booster_RESOURCE_TYPE, (void *)&(booster_resource))) { ret = exg_error(env, "Invalid Booster"); goto END; } - if (!exg_get_string(env, argv[1], &indptr)) { - ret = exg_error(env, "Indptr must be a JSON-encoded string"); + + // Build ArrayInterface JSON for each sparse array from components + if (!exg_build_array_interface_json(env, argv[1], argv[2], argv[3], argv[4], &indptr, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract indptr ArrayInterface"); goto END; } - if (!exg_get_string(env, argv[2], &indices)) { - ret = exg_error(env, "Indices must be a JSON-encoded string"); + + if (!exg_build_array_interface_json(env, argv[5], argv[6], argv[7], argv[8], &indices, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract indices ArrayInterface"); goto END; } - if (!exg_get_string(env, argv[3], &data)) { - ret = exg_error(env, "Data must be a JSON-encoded string"); + + if (!exg_build_array_interface_json(env, argv[9], argv[10], argv[11], argv[12], &data, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract data ArrayInterface"); goto END; } - if (!enif_get_int(env, argv[4], &ncols)) { + + if (!enif_get_int(env, argv[13], &ncols)) { ret = exg_error(env, "Ncols must be an integer"); goto END; } - if (!exg_get_string(env, argv[5], &config)) { + + if (ncols < 0) { + ret = exg_error(env, "Ncols must be non-negative"); + goto END; + } + + if (!exg_get_string(env, argv[14], &config)) { ret = exg_error(env, "Config must be a JSON-encoded string"); goto END; } - if (!enif_get_resource(env, argv[6], DMatrix_RESOURCE_TYPE, + + if (!enif_get_resource(env, argv[15], DMatrix_RESOURCE_TYPE, (void *)&(proxy_resource))) { proxy = NULL; } else { @@ -772,8 +835,9 @@ ERL_NIF_TERM EXGBoosterPredictFromCSR(ErlNifEnv *env, int argc, } booster = *booster_resource; result = - XGBoosterPredictFromCSR(booster, indptr, indices, data, ncols, config, - proxy, &out_shape, &out_dim, &out_result); + XGBoosterPredictFromCSR(booster, indptr, indices, data, + (bst_ulong)ncols, config, proxy, &out_shape, + &out_dim, &out_result); if (result == 0) { ret = collect_prediction_results(env, out_shape, out_dim, out_result); } else { @@ -818,6 +882,7 @@ ERL_NIF_TERM EXGBoosterLoadModel(ErlNifEnv *env, int argc, if (result == 0) { ret = make_Booster_resource(env, booster); } else { + XGBoosterFree(booster); ret = exg_error(env, XGBGetLastError()); } END: @@ -866,7 +931,7 @@ ERL_NIF_TERM EXGBoosterSerializeToBuffer(ErlNifEnv *env, int argc, BoosterHandle booster; BoosterHandle **booster_resource = NULL; bst_ulong out_len = 0; - char *out_buf = NULL; + const char *out_buf = NULL; int result = -1; ERL_NIF_TERM ret = -1; ErlNifBinary out_bin; @@ -894,6 +959,7 @@ ERL_NIF_TERM EXGBoosterSerializeToBuffer(ErlNifEnv *env, int argc, END: return ret; } + ERL_NIF_TERM EXGBoosterDeserializeFromBuffer(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { BoosterHandle booster; @@ -920,6 +986,7 @@ ERL_NIF_TERM EXGBoosterDeserializeFromBuffer(ErlNifEnv *env, int argc, if (result == 0) { ret = make_Booster_resource(env, booster); } else { + XGBoosterFree(booster); ret = exg_error(env, XGBGetLastError()); } END: @@ -955,6 +1022,7 @@ ERL_NIF_TERM EXGBoosterLoadModelFromBuffer(ErlNifEnv *env, int argc, if (result == 0) { ret = make_Booster_resource(env, booster); } else { + XGBoosterFree(booster); ret = exg_error(env, XGBGetLastError()); } END: @@ -969,7 +1037,7 @@ ERL_NIF_TERM EXGBoosterSaveModelToBuffer(ErlNifEnv *env, int argc, BoosterHandle booster; BoosterHandle **booster_resource = NULL; bst_ulong out_len = 0; - char *out_buf = NULL; + const char *out_buf = NULL; char *config = NULL; int result = -1; ERL_NIF_TERM ret = -1; @@ -1012,7 +1080,7 @@ ERL_NIF_TERM EXGBoosterSaveJsonConfig(ErlNifEnv *env, int argc, BoosterHandle booster; BoosterHandle **booster_resource = NULL; bst_ulong out_len = 0; - char *out_buf = NULL; + const char *out_buf = NULL; int result = -1; ERL_NIF_TERM ret = -1; ErlNifBinary out_bin; @@ -1083,7 +1151,7 @@ ERL_NIF_TERM EXGBoosterDumpModelEx(ErlNifEnv *env, int argc, BoosterHandle booster; BoosterHandle **booster_resource = NULL; bst_ulong out_len = 0; - char **out_dump_array = NULL; + const char **out_dump_array = NULL; char *fmap = NULL; int with_stats = 0; char *format = NULL; @@ -1116,10 +1184,7 @@ ERL_NIF_TERM EXGBoosterDumpModelEx(ErlNifEnv *env, int argc, if (result == 0) { ERL_NIF_TERM arr[out_len]; for (bst_ulong i = 0; i < out_len; ++i) { - char *local = enif_alloc(strlen(out_dump_array[i]) + 1); - strcpy(local, out_dump_array[i]); - arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1); - // TODO: Do we free here or is it handled by the XGBoost library / BEAM ? + arr[i] = enif_make_string(env, out_dump_array[i], ERL_NIF_LATIN1); } ret = exg_ok(env, enif_make_list_from_array(env, arr, out_len)); } else { diff --git a/c/exgboost/src/config.c b/c/exgboost/src/config.c index 89f432e..9590d66 100644 --- a/c/exgboost/src/config.c +++ b/c/exgboost/src/config.c @@ -59,7 +59,7 @@ ERL_NIF_TERM EXGBSetGlobalConfig(ErlNifEnv *env, int argc, ERL_NIF_TERM EXGBGetGlobalConfig(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { - char *out = NULL; + const char *out = NULL; int result = -1; ERL_NIF_TERM ret = 0; if (argc != 0) { @@ -68,7 +68,7 @@ ERL_NIF_TERM EXGBGetGlobalConfig(ErlNifEnv *env, int argc, } // No need to free out, it's a pointer to a static string defined in the // xgboost config struct - result = XGBGetGlobalConfig((char const **)&out); + result = XGBGetGlobalConfig(&out); if (result == 0) { ret = exg_ok(env, enif_make_string(env, out, ERL_NIF_LATIN1)); } else { @@ -76,4 +76,4 @@ ERL_NIF_TERM EXGBGetGlobalConfig(ErlNifEnv *env, int argc, } END: return ret; -} \ No newline at end of file +} diff --git a/c/exgboost/src/dmatrix.c b/c/exgboost/src/dmatrix.c index 75dc621..6e3a551 100644 --- a/c/exgboost/src/dmatrix.c +++ b/c/exgboost/src/dmatrix.c @@ -1,4 +1,6 @@ #include "dmatrix.h" +#include +#include static ERL_NIF_TERM make_DMatrix_resource(ErlNifEnv *env, DMatrixHandle handle) { @@ -7,6 +9,7 @@ static ERL_NIF_TERM make_DMatrix_resource(ErlNifEnv *env, enif_alloc_resource(DMatrix_RESOURCE_TYPE, sizeof(DMatrixHandle *)); if (resource != NULL) { *resource = handle; + // BEAM resource now owns the handle and releases it in resource cleanup. ret = exg_ok(env, enif_make_resource(env, resource)); enif_release_resource(resource); } else { @@ -138,38 +141,45 @@ ERL_NIF_TERM EXGDMatrixCreateFromSparse(ErlNifEnv *env, int argc, char *indptr_interface = NULL; char *indices_interface = NULL; char *data_interface = NULL; + const char *error_msg = NULL; int n = 0; char *config = NULL; char *format = NULL; DMatrixHandle handle; ERL_NIF_TERM ret = 0; - if (argc != 6) { + + if (argc != 15) { ret = exg_error(env, "Wrong number of arguments"); goto END; } - if (!exg_get_string(env, argv[0], &indptr_interface)) { - ret = - exg_error(env, "Indptr Array Interface must be a JSON-Encoded string"); + + // Extract ArrayInterfaces with fresh addresses: (binary, typestr, shape, readonly) per array + if (!exg_build_array_interface_json(env, argv[0], argv[1], argv[2], argv[3], &indptr_interface, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract indptr ArrayInterface"); goto END; } - if (!exg_get_string(env, argv[1], &indices_interface)) { - ret = - exg_error(env, "Indices Array Interface must be a JSON-Encoded string"); + + if (!exg_build_array_interface_json(env, argv[4], argv[5], argv[6], argv[7], &indices_interface, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract indices ArrayInterface"); goto END; } - if (!exg_get_string(env, argv[2], &data_interface)) { - ret = exg_error(env, "Data Array Interface must be a JSON-Encoded string"); + + if (!exg_build_array_interface_json(env, argv[8], argv[9], argv[10], argv[11], &data_interface, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract data ArrayInterface"); goto END; } - if (!enif_get_int(env, argv[3], &n)) { + + if (!enif_get_int(env, argv[12], &n)) { ret = exg_error(env, "Ncol must be an integer"); goto END; } - if (!exg_get_string(env, argv[4], &config)) { + + if (!exg_get_string(env, argv[13], &config)) { ret = exg_error(env, "Config must be a string"); goto END; } - if (!exg_get_string(env, argv[5], &format)) { + + if (!exg_get_string(env, argv[14], &format)) { ret = exg_error(env, "Format must be a string"); goto END; } @@ -212,86 +222,27 @@ ERL_NIF_TERM EXGDMatrixCreateFromSparse(ErlNifEnv *env, int argc, return ret; } -ERL_NIF_TERM EXGDMatrixCreateFromCSREx(ErlNifEnv *env, int argc, - const ERL_NIF_TERM argv[]) { - ErlNifBinary indptr_bin; - ErlNifBinary indices_bin; - ErlNifBinary data_bin; - int result = -1; - ErlNifUInt64 *indptr = NULL; - uint32_t *indices = NULL; - float *data = NULL; - uint32_t nindptr = 0; - ErlNifUInt64 nelem = 0; - ErlNifUInt64 ncol = 0; - DMatrixHandle handle; - ERL_NIF_TERM ret = 0; - if (argc != 6) { - ret = exg_error(env, "Wrong number of arguments"); - goto END; - } - if (!enif_inspect_binary(env, argv[0], &indptr_bin)) { - ret = exg_error(env, "Indptr must be a binary of uint64_t"); - goto END; - } - if (!enif_inspect_binary(env, argv[1], &indices_bin)) { - ret = exg_error(env, "Indices must be a binary of uint64_t"); - goto END; - } - if (!enif_inspect_binary(env, argv[2], &data_bin)) { - ret = exg_error(env, "Data must be a binary of uint64_t"); - goto END; - } - if (!enif_get_uint(env, argv[3], &nindptr)) { - ret = exg_error(env, "Nindptr must be a uint64_t"); - goto END; - } - if (!enif_get_uint64(env, argv[4], &nelem)) { - ret = exg_error(env, "Nelem must be a uint64_t"); - goto END; - } - if (!enif_get_uint64(env, argv[5], &ncol)) { - ret = exg_error(env, "Ncol must be a uint64_t"); - goto END; - } - indptr = (ErlNifUInt64 *)indptr_bin.data; - indices = (uint32_t *)indices_bin.data; - data = (float *)data_bin.data; - if (indptr_bin.size != nindptr * sizeof(ErlNifUInt64)) { - ret = exg_error(env, "Indptr size does not match nindptr"); - goto END; - } - if (data_bin.size != nelem * sizeof(float)) { - ret = exg_error(env, "Data size does not match nelem"); - goto END; - } - result = XGDMatrixCreateFromCSREx(indptr, indices, data, nindptr, nelem, ncol, - &handle); - if (result == 0) { - ret = exg_ok(env, enif_make_resource(env, handle)); - enif_release_resource(handle); - } else { - ret = exg_error(env, XGBGetLastError()); - } -END: - return ret; -} - ERL_NIF_TERM EXGDMatrixCreateFromDense(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { int result = -1; char *array_interface = NULL; + const char *error_msg = NULL; char *config = NULL; DMatrixHandle out; ERL_NIF_TERM ret = 0; - if (argc != 2) { + + if (argc != 5) { ret = exg_error(env, "Wrong number of arguments"); + goto END; } - if (!exg_get_string(env, argv[0], &array_interface)) { - ret = exg_error(env, "Array Interface must be a JSON-Encoded string"); + + // Build ArrayInterface JSON from components: (binary, typestr, shape, readonly) + if (!exg_build_array_interface_json(env, argv[0], argv[1], argv[2], argv[3], &array_interface, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract ArrayInterface"); goto END; } - if (!exg_get_string(env, argv[1], &config)) { + + if (!exg_get_string(env, argv[4], &config)) { ret = exg_error(env, "Config must be a JSON-Encoded string"); goto END; } @@ -343,16 +294,21 @@ ERL_NIF_TERM EXGDMatrixSetStrFeatureInfo(ErlNifEnv *env, int argc, goto END; } handle = *resource; - result = XGDMatrixSetStrFeatureInfo(handle, field, features, num_features); + // XGBoost reads features during the call; caller keeps ownership. + result = XGDMatrixSetStrFeatureInfo(handle, field, + (const char **)features, num_features); if (result == 0) { ret = ok_atom(env); } else { ret = exg_error(env, XGBGetLastError()); } END: + // Helper-allocated buffers must be reclaimed on all paths. if (features != NULL) { - enif_free(features); - features = NULL; + exg_free_string_list(features, num_features); + } + if (field != NULL) { + enif_free(field); } return ret; } @@ -388,14 +344,17 @@ ERL_NIF_TERM EXGDMatrixGetStrFeatureInfo(ErlNifEnv *env, int argc, result = XGDMatrixGetStrFeatureInfo(handle, field, &out_size, &c_out_features); if (result == 0) { + // Check VLA size fits in size_t + if (out_size > SIZE_MAX / sizeof(ERL_NIF_TERM)) { + ret = exg_error(env, "Result is too large"); + goto END; + } ERL_NIF_TERM arr[out_size]; for (bst_ulong i = 0; i < out_size; ++i) { - char *local = enif_alloc(strlen(c_out_features[i]) + 1); - strcpy(local, c_out_features[i]); - arr[i] = enif_make_string(env, local, ERL_NIF_LATIN1); - // TODO: Do we free here or is it handled by the XGBoost library / BEAM? + // enif_make_string materializes a BEAM term; no temporary C copy needed. + arr[i] = enif_make_string(env, c_out_features[i], ERL_NIF_LATIN1); } - ret = exg_ok(env, enif_make_list_from_array(env, arr, out_size)); + ret = exg_ok(env, enif_make_list_from_array(env, arr, (size_t)out_size)); } else { ret = exg_error(env, XGBGetLastError()); } @@ -433,10 +392,12 @@ ERL_NIF_TERM EXGDMatrixSetDenseInfo(ErlNifEnv *env, int argc, ret = exg_error(env, "Data must be a binary"); goto END; } - if (!enif_get_ulong(env, argv[3], &size)) { - ret = exg_error(env, "Size must be an integer"); + ErlNifUInt64 size_arg; + if (!enif_get_uint64(env, argv[3], &size_arg)) { + ret = exg_error(env, "Size must be a non-negative integer"); goto END; } + size = (bst_ulong)size_arg; if (!enif_get_int(env, argv[4], &type)) { ret = exg_error(env, "Type must be an integer"); goto END; @@ -488,7 +449,7 @@ ERL_NIF_TERM EXGDMatrixNumRow(ErlNifEnv *env, int argc, handle = *resource; result = XGDMatrixNumRow(handle, &out); if (result == 0) { - ret = exg_ok(env, enif_make_ulong(env, out)); + ret = exg_ok(env, enif_make_uint64(env, (ErlNifUInt64)out)); } else { ret = exg_error(env, XGBGetLastError()); } @@ -515,7 +476,7 @@ ERL_NIF_TERM EXGDMatrixNumCol(ErlNifEnv *env, int argc, handle = *resource; result = XGDMatrixNumCol(handle, &out); if (result == 0) { - ret = exg_ok(env, enif_make_ulong(env, out)); + ret = exg_ok(env, enif_make_uint64(env, (ErlNifUInt64)out)); } else { ret = exg_error(env, XGBGetLastError()); } @@ -542,7 +503,7 @@ ERL_NIF_TERM EXGDMatrixNumNonMissing(ErlNifEnv *env, int argc, handle = *resource; result = XGDMatrixNumNonMissing(handle, &out); if (result == 0) { - ret = exg_ok(env, enif_make_ulong(env, out)); + ret = exg_ok(env, enif_make_uint64(env, (ErlNifUInt64)out)); } else { ret = exg_error(env, XGBGetLastError()); } @@ -556,25 +517,32 @@ ERL_NIF_TERM EXGDMatrixSetInfoFromInterface(ErlNifEnv *env, int argc, DMatrixHandle **resource = NULL; char *field = NULL; char *data_interface = NULL; + const char *error_msg = NULL; int result = -1; ERL_NIF_TERM ret = 0; - if (argc != 3) { + + if (argc != 6) { ret = exg_error(env, "Wrong number of arguments"); goto END; } + if (!enif_get_resource(env, argv[0], DMatrix_RESOURCE_TYPE, (void *)&resource)) { ret = exg_error(env, "DMatrix must be a resource"); goto END; } + if (!exg_get_string(env, argv[1], &field)) { ret = exg_error(env, "Field must be a string"); goto END; } - if (!exg_get_string(env, argv[2], &data_interface)) { - ret = exg_error(env, "Data must be a string"); + + // Build ArrayInterface JSON from components: (binary, typestr, shape, readonly) + if (!exg_build_array_interface_json(env, argv[2], argv[3], argv[4], argv[5], &data_interface, &error_msg)) { + ret = exg_error(env, error_msg ? error_msg : "Failed to extract data ArrayInterface"); goto END; } + if (strcmp(field, "label") != 0 && strcmp(field, "weight") != 0 && strcmp(field, "base_margin") != 0 && strcmp(field, "group") != 0 && strcmp(field, "label_lower_bound") != 0 && @@ -646,7 +614,7 @@ ERL_NIF_TERM EXGDMatrixGetFloatInfo(ErlNifEnv *env, int argc, DMatrixHandle handle; DMatrixHandle **resource = NULL; char *field = NULL; - float *out = NULL; + const float *out = NULL; bst_ulong len = 0; int result = -1; ERL_NIF_TERM ret = 0; @@ -679,6 +647,7 @@ ERL_NIF_TERM EXGDMatrixGetFloatInfo(ErlNifEnv *env, int argc, if (result == 0) { arr = enif_alloc(sizeof(ERL_NIF_TERM) * len); for (int i = 0; i < len; i++) { + // Values are copied into BEAM-managed terms here. arr[i] = enif_make_double(env, out[i]); } ret = exg_ok(env, enif_make_list_from_array(env, arr, len)); @@ -700,7 +669,7 @@ ERL_NIF_TERM EXGDMatrixGetUIntInfo(ErlNifEnv *env, int argc, DMatrixHandle handle; DMatrixHandle **resource = NULL; char *field = NULL; - unsigned *out = NULL; + const unsigned *out = NULL; bst_ulong len = 0; int result = -1; ERL_NIF_TERM ret = 0; @@ -728,6 +697,7 @@ ERL_NIF_TERM EXGDMatrixGetUIntInfo(ErlNifEnv *env, int argc, if (result == 0) { arr = enif_alloc(sizeof(ERL_NIF_TERM) * len); for (int i = 0; i < len; i++) { + // Values are copied into BEAM-managed terms here. arr[i] = enif_make_uint(env, out[i]); } ret = exg_ok(env, enif_make_list_from_array(env, arr, len)); @@ -783,9 +753,16 @@ ERL_NIF_TERM EXGDMatrixGetDataAsCSR(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); goto END; } - out_indptr = malloc(sizeof(bst_ulong) * (num_rows + 1)); - out_indices = malloc(sizeof(unsigned) * num_non_missing); - out_data = malloc(sizeof(float) * num_non_missing); + // Check allocation sizes fit in size_t + if (num_rows > SIZE_MAX / sizeof(bst_ulong) - 1 || + num_non_missing > SIZE_MAX / sizeof(unsigned) || + num_non_missing > SIZE_MAX / sizeof(float)) { + ret = exg_error(env, "Matrix is too large"); + goto END; + } + out_indptr = malloc(sizeof(bst_ulong) * (size_t)(num_rows + 1)); + out_indices = malloc(sizeof(unsigned) * (size_t)num_non_missing); + out_data = malloc(sizeof(float) * (size_t)num_non_missing); if (!out_indptr || !out_indices || !out_data) { ret = exg_error(env, "Failed to allocate memory"); goto END; @@ -796,26 +773,33 @@ ERL_NIF_TERM EXGDMatrixGetDataAsCSR(ErlNifEnv *env, int argc, ret = exg_error(env, XGBGetLastError()); goto END; } - indptr = enif_alloc(sizeof(ERL_NIF_TERM) * (num_rows + 1)); - indices = enif_alloc(sizeof(ERL_NIF_TERM) * num_non_missing); - data = enif_alloc(sizeof(ERL_NIF_TERM) * num_non_missing); + // Check enif_alloc sizes fit in size_t + if (num_rows > SIZE_MAX / sizeof(ERL_NIF_TERM) - 1 || + num_non_missing > SIZE_MAX / sizeof(ERL_NIF_TERM)) { + ret = exg_error(env, "Matrix is too large"); + goto END; + } + indptr = enif_alloc(sizeof(ERL_NIF_TERM) * (size_t)(num_rows + 1)); + indices = enif_alloc(sizeof(ERL_NIF_TERM) * (size_t)num_non_missing); + data = enif_alloc(sizeof(ERL_NIF_TERM) * (size_t)num_non_missing); if (!indptr || !indices || !data) { ret = exg_error(env, "Failed to allocate memory"); goto END; } - for (int i = 0; i < num_rows + 1; i++) { - indptr[i] = enif_make_ulong(env, out_indptr[i]); + for (bst_ulong i = 0; i < num_rows + 1; ++i) { + indptr[i] = enif_make_uint64(env, (ErlNifUInt64)out_indptr[i]); } - for (int i = 0; i < num_non_missing; i++) { + for (bst_ulong i = 0; i < num_non_missing; ++i) { indices[i] = enif_make_uint(env, out_indices[i]); data[i] = enif_make_double(env, out_data[i]); } ret = exg_ok(env, enif_make_tuple3( - env, enif_make_list_from_array(env, indptr, num_rows + 1), - enif_make_list_from_array(env, indices, num_non_missing), - enif_make_list_from_array(env, data, num_non_missing))); + env, enif_make_list_from_array(env, indptr, (size_t)(num_rows + 1)), + enif_make_list_from_array(env, indices, (size_t)num_non_missing), + enif_make_list_from_array(env, data, (size_t)num_non_missing))); END: + // Mixed allocators: enif_free for enif_alloc buffers, free for malloc buffers. if (config != NULL) { enif_free(config); config = NULL; @@ -832,6 +816,15 @@ ERL_NIF_TERM EXGDMatrixGetDataAsCSR(ErlNifEnv *env, int argc, free(out_data); out_data = NULL; } + if (indptr != NULL) { + enif_free(indptr); + } + if (indices != NULL) { + enif_free(indices); + } + if (data != NULL) { + enif_free(data); + } return ret; }; @@ -911,13 +904,98 @@ ERL_NIF_TERM EXGProxyDMatrixCreate(ErlNifEnv *env, int argc, return ret; } +// Helper to parse ArrayInterface JSON and copy data atomically +static int exg_parse_and_copy_array_interface( + ErlNifEnv *env, + const char *json_str, + ERL_NIF_TERM *out_map +) { + // Parse the JSON to extract address, typestr, shape + // For now, we'll use a simple parser since the format is predictable: + // {"version":3,"typestr":" e.g. " SIZE_MAX / bytes_per_elem) { + return 0; + } + size_t total_size = shape_len * bytes_per_elem; + + // Copy data from XGBoost-owned address immediately + ErlNifBinary data_bin; + if (address == 0 && total_size != 0) { + return 0; + } + if (!enif_alloc_binary(total_size, &data_bin)) { + return 0; + } + if (total_size != 0) { + memcpy(data_bin.data, (const void *)address, total_size); + } + + // Build result map: %{binary: binary, typestr: string, shape: [n]} + ERL_NIF_TERM binary_key = enif_make_atom(env, "binary"); + ERL_NIF_TERM typestr_key = enif_make_atom(env, "typestr"); + ERL_NIF_TERM shape_key = enif_make_atom(env, "shape"); + + // Convert typestr to binary string + size_t typestr_len = strlen(typestr); + ErlNifBinary typestr_bin; + if (!enif_alloc_binary(typestr_len, &typestr_bin)) { + return 0; + } + memcpy(typestr_bin.data, typestr, typestr_len); + + ERL_NIF_TERM keys[] = {binary_key, typestr_key, shape_key}; + ERL_NIF_TERM values[] = { + enif_make_binary(env, &data_bin), + enif_make_binary(env, &typestr_bin), + enif_make_list1(env, enif_make_uint64(env, (ErlNifUInt64)shape_len)) + }; + + return enif_make_map_from_arrays(env, keys, values, 3, out_map); +} + ERL_NIF_TERM EXGDMatrixGetQuantileCut(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { DMatrixHandle handle; DMatrixHandle **resource = NULL; char *config = NULL; - char const *out_indptr = NULL; - char const *out_data = NULL; + char const *out_indptr_json = NULL; + char const *out_data_json = NULL; + ERL_NIF_TERM indptr_map, data_map; ERL_NIF_TERM ret = -1; int result = -1; @@ -925,8 +1003,7 @@ ERL_NIF_TERM EXGDMatrixGetQuantileCut(ErlNifEnv *env, int argc, ret = exg_error(env, "Wrong number of arguments"); goto END; } - if (!enif_get_resource(env, argv[0], DMatrix_RESOURCE_TYPE, - (void *)&resource)) { + if (!enif_get_resource(env, argv[0], DMatrix_RESOURCE_TYPE, (void *)&resource)) { ret = exg_error(env, "DMatrix must be a resource"); goto END; } @@ -935,15 +1012,26 @@ ERL_NIF_TERM EXGDMatrixGetQuantileCut(ErlNifEnv *env, int argc, goto END; } handle = *resource; - result = XGDMatrixGetQuantileCut(handle, config, &out_indptr, &out_data); - if (result == 0) { - ret = exg_ok( - env, - enif_make_tuple2(env, enif_make_string(env, out_indptr, ERL_NIF_LATIN1), - enif_make_string(env, out_data, ERL_NIF_LATIN1))); - } else { + + // Get ArrayInterface JSON from XGBoost + result = XGDMatrixGetQuantileCut(handle, config, &out_indptr_json, &out_data_json); + if (result != 0) { ret = exg_error(env, XGBGetLastError()); + goto END; } + + // Parse JSON and copy data atomically, before any other operations + if (!exg_parse_and_copy_array_interface(env, out_indptr_json, &indptr_map)) { + ret = exg_error(env, "Failed to parse indptr ArrayInterface"); + goto END; + } + if (!exg_parse_and_copy_array_interface(env, out_data_json, &data_map)) { + ret = exg_error(env, "Failed to parse data ArrayInterface"); + goto END; + } + + ret = exg_ok(env, enif_make_tuple2(env, indptr_map, data_map)); + END: if (config != NULL) { enif_free(config); diff --git a/c/exgboost/src/exgboost.c b/c/exgboost/src/exgboost.c index 96877ae..ab5e82f 100644 --- a/c/exgboost/src/exgboost.c +++ b/c/exgboost/src/exgboost.c @@ -1,6 +1,8 @@ #include "exgboost.h" static int load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info) { + exg_init_atoms(env); + DMatrix_RESOURCE_TYPE = enif_open_resource_type( env, NULL, "DMatrix_RESOURCE_TYPE", DMatrix_RESOURCE_TYPE_cleanup, (ErlNifResourceFlags)(ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER), NULL); @@ -15,6 +17,8 @@ static int load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info) { static int upgrade(ErlNifEnv *env, void **priv_data, void **old_priv_data, ERL_NIF_TERM load_info) { + exg_init_atoms(env); + DMatrix_RESOURCE_TYPE = enif_open_resource_type( env, NULL, "DMatrix_RESOURCE_TYPE", DMatrix_RESOURCE_TYPE_cleanup, ERL_NIF_RT_TAKEOVER, NULL); @@ -39,17 +43,21 @@ static ErlNifFunc nif_funcs[] = { {"dmatrix_create_from_uri", 1, EXGDMatrixCreateFromURI, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"dmatrix_create_from_mat", 4, EXGDMatrixCreateFromMat}, - {"dmatrix_create_from_sparse", 6, EXGDMatrixCreateFromSparse}, - {"dmatrix_create_from_dense", 2, EXGDMatrixCreateFromDense}, + {"dmatrix_create_from_sparse", 15, EXGDMatrixCreateFromSparse}, + {"dmatrix_create_from_dense", 5, EXGDMatrixCreateFromDense}, {"dmatrix_set_str_feature_info", 3, EXGDMatrixSetStrFeatureInfo}, {"dmatrix_get_str_feature_info", 2, EXGDMatrixGetStrFeatureInfo}, {"dmatrix_num_row", 1, EXGDMatrixNumRow}, {"dmatrix_num_col", 1, EXGDMatrixNumCol}, {"dmatrix_num_non_missing", 1, EXGDMatrixNumNonMissing}, - {"dmatrix_set_info_from_interface", 3, EXGDMatrixSetInfoFromInterface}, + {"dmatrix_set_info_from_interface", 6, EXGDMatrixSetInfoFromInterface}, {"dmatrix_save_binary", 3, EXGDMatrixSaveBinary}, + + // Unsafe, to be deprecated in future releases. {"get_binary_address", 1, exg_get_binary_address}, + // Unsafe, to be deprecated in future releases. {"get_binary_from_address", 2, exg_get_binary_from_address}, + {"dmatrix_get_float_info", 2, EXGDMatrixGetFloatInfo}, {"dmatrix_get_uint_info", 2, EXGDMatrixGetUIntInfo}, {"dmatrix_get_data_as_csr", 2, EXGDMatrixGetDataAsCSR}, @@ -74,9 +82,9 @@ static ErlNifFunc nif_funcs[] = { {"booster_slice", 4, EXGBoosterSlice}, {"booster_predict_from_dmatrix", 3, EXGBoosterPredictFromDMatrix, ERL_NIF_DIRTY_JOB_CPU_BOUND}, - {"booster_predict_from_dense", 4, EXGBoosterPredictFromDense, + {"booster_predict_from_dense", 7, EXGBoosterPredictFromDense, ERL_NIF_DIRTY_JOB_CPU_BOUND}, - {"booster_predict_from_csr", 7, EXGBoosterPredictFromCSR, + {"booster_predict_from_csr", 16, EXGBoosterPredictFromCSR, ERL_NIF_DIRTY_JOB_CPU_BOUND}, {"booster_load_model", 1, EXGBoosterLoadModel, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"booster_save_model", 2, EXGBoosterSaveModel, ERL_NIF_DIRTY_JOB_IO_BOUND}, diff --git a/c/exgboost/src/utils.c b/c/exgboost/src/utils.c index 401def7..f893706 100644 --- a/c/exgboost/src/utils.c +++ b/c/exgboost/src/utils.c @@ -1,4 +1,23 @@ #include "utils.h" +#include +#include +#include + +// Ensure bst_ulong and ErlNifUInt64 are both 64-bit for safe conversions +_Static_assert( + sizeof(bst_ulong) == sizeof(ErlNifUInt64), + "bst_ulong and ErlNifUInt64 must both be 64-bit" +); + +// Cached atoms +static ERL_NIF_TERM ATOM_TRUE; +static ERL_NIF_TERM ATOM_FALSE; + +// Initialize atoms (call this from NIF load) +void exg_init_atoms(ErlNifEnv *env) { + ATOM_TRUE = enif_make_atom(env, "true"); + ATOM_FALSE = enif_make_atom(env, "false"); +} // Atoms ERL_NIF_TERM exg_error(ErlNifEnv *env, const char *msg) { @@ -61,12 +80,14 @@ int exg_get_list(ErlNifEnv *env, ERL_NIF_TERM term, double **out) { return 0; } *out = (double *)enif_alloc(len * sizeof(double)); - if (out == NULL) { + if (*out == NULL) { return 0; } while (enif_get_list_cell(env, term, &head, &tail)) { int ret = enif_get_double(env, head, &((*out)[i])); if (!ret) { + enif_free(*out); + *out = NULL; return 0; } term = tail; @@ -89,6 +110,8 @@ int exg_get_string_list(ErlNifEnv *env, ERL_NIF_TERM term, char ***out, while (enif_get_list_cell(env, term, &head, &tail)) { int ret = exg_get_string(env, head, &((*out)[i])); if (!ret) { + exg_free_string_list(*out, i); + *out = NULL; return 0; } term = tail; @@ -105,13 +128,15 @@ int exg_get_dmatrix_list(ErlNifEnv *env, ERL_NIF_TERM term, return 0; } *dmats = (DMatrixHandle *)enif_alloc(*len * sizeof(DMatrixHandle)); - if (NULL == dmats) { + if (NULL == *dmats) { return 0; } while (enif_get_list_cell(env, term, &head, &tail)) { DMatrixHandle **resource = NULL; if (!enif_get_resource(env, head, DMatrix_RESOURCE_TYPE, (void *)&(resource))) { + exg_free_dmatrix_list(*dmats); + *dmats = NULL; return 0; } memcpy(&((*dmats)[i]), resource, sizeof(DMatrixHandle)); @@ -121,6 +146,25 @@ int exg_get_dmatrix_list(ErlNifEnv *env, ERL_NIF_TERM term, return 1; } +void exg_free_string_list(char **items, unsigned len) { + if (items == NULL) { + return; + } + for (unsigned i = 0; i < len; ++i) { + if (items[i] != NULL) { + enif_free(items[i]); + } + } + enif_free(items); +} + +void exg_free_dmatrix_list(DMatrixHandle *dmats) { + if (dmats != NULL) { + enif_free(dmats); + } +} + +// Unsafe, to be deprecated in future releases. ERL_NIF_TERM exg_get_binary_address(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { ErlNifBinary bin; @@ -138,6 +182,7 @@ ERL_NIF_TERM exg_get_binary_address(ErlNifEnv *env, int argc, return ret; } +// Unsafe, to be deprecated in future releases. ERL_NIF_TERM exg_get_binary_from_address(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { ErlNifBinary out_bin; @@ -166,6 +211,299 @@ ERL_NIF_TERM exg_get_binary_from_address(ErlNifEnv *env, int argc, return ret; } +// Helper: Extract boolean from term (true/false atom only) +static int exg_get_boolean(ErlNifEnv *env, ERL_NIF_TERM term, int *value) { + if (enif_is_identical(term, ATOM_TRUE)) { + *value = 1; + return 1; + } else if (enif_is_identical(term, ATOM_FALSE)) { + *value = 0; + return 1; + } + return 0; +} + +// Helper: Validate typestr format +// Accepts patterns like: = capacity) { // +2 for ']' and '\0' + capacity *= 2; + char *new_json = enif_realloc(json, capacity); + if (new_json == NULL) { + enif_free(json); + *error_msg = "Failed to reallocate memory for shape JSON"; + return 0; + } + json = new_json; + } + + // Append dimension + strcpy(json + pos, dim_str); + pos += written; + + // Add comma if not last + if (i < shape_len - 1) { + json[pos++] = ','; + } + } + + json[pos++] = ']'; + json[pos] = '\0'; + + *json_out = json; + return 1; +} + +// Helper: Validate shape and check binary size +static int exg_validate_shape_and_size(ErlNifEnv *env, ERL_NIF_TERM shape_term, + const char *typestr, size_t binary_size, + size_t *required_bytes_out, + const char **error_msg) { + unsigned shape_len; + size_t element_size = exg_bytes_per_element(typestr); + + if (element_size == 0) { + *error_msg = "Invalid typestr: element size is 0"; + return 0; + } + + if (!enif_get_list_length(env, shape_term, &shape_len)) { + *error_msg = "Shape must be a proper list"; + return 0; + } + + // Calculate total elements with overflow checking + size_t total_elements = 1; + ERL_NIF_TERM head, tail = shape_term; + + for (unsigned i = 0; i < shape_len; i++) { + if (!enif_get_list_cell(env, tail, &head, &tail)) { + *error_msg = "Failed to iterate shape list"; + return 0; + } + + ErlNifUInt64 dim; + if (!enif_get_uint64(env, head, &dim)) { + *error_msg = "Shape values must be non-negative integers"; + return 0; + } + + // Check for overflow in multiplication + if (dim > 0 && total_elements > SIZE_MAX / dim) { + *error_msg = "Shape dimensions overflow"; + return 0; + } + + total_elements *= (size_t)dim; + } + + // Check for overflow in byte calculation + if (total_elements > SIZE_MAX / element_size) { + *error_msg = "Required byte size overflows"; + return 0; + } + + size_t required_bytes = total_elements * element_size; + + if (binary_size < required_bytes) { + *error_msg = "Binary is too small for the specified shape and type"; + return 0; + } + + *required_bytes_out = required_bytes; + return 1; +} + +// Helper to build Array Interface JSON from components with fresh address +// Returns 1 on success, 0 on failure +// Helper to build Array Interface JSON from components with fresh address +// Returns 1 on success, 0 on failure +int exg_build_array_interface_json(ErlNifEnv *env, ERL_NIF_TERM binary_term, + ERL_NIF_TERM typestr_term, + ERL_NIF_TERM shape_term, + ERL_NIF_TERM readonly_term, char **json_out, + const char **error_msg) { + ErlNifBinary data_bin; + char *typestr = NULL; + char *shape_json = NULL; + char *json = NULL; + int readonly; + int ok = 0; + + *json_out = NULL; + *error_msg = NULL; + + // Extract and validate binary + if (!enif_inspect_binary(env, binary_term, &data_bin)) { + *error_msg = "Binary argument required"; + goto CLEANUP; + } + + // Extract typestr + if (!exg_get_string(env, typestr_term, &typestr)) { + *error_msg = "Typestr must be a string"; + goto CLEANUP; + } + + // Validate typestr format + if (!exg_valid_typestr(typestr)) { + *error_msg = "Unsupported typestr format"; + goto CLEANUP; + } + + // Extract and validate readonly + if (!exg_get_boolean(env, readonly_term, &readonly)) { + *error_msg = "Readonly must be a boolean (true or false atom)"; + goto CLEANUP; + } + + // Validate shape dimensions and binary size + size_t required_bytes; + if (!exg_validate_shape_and_size(env, shape_term, typestr, data_bin.size, + &required_bytes, error_msg)) { + goto CLEANUP; + } + + // Build shape JSON + if (!exg_shape_to_json(env, shape_term, &shape_json, error_msg)) { + goto CLEANUP; + } + + // Get binary address + uintptr_t address = (uintptr_t)data_bin.data; + + // Calculate required JSON buffer size + int needed = snprintf(NULL, 0, + "{\"typestr\":\"%s\",\"shape\":%s," + "\"data\":[%" PRIuPTR ",%s],\"version\":3}", + typestr, shape_json, address, + readonly ? "true" : "false"); + + if (needed < 0) { + *error_msg = "Failed to calculate JSON size"; + goto CLEANUP; + } + + // Allocate JSON buffer + json = enif_alloc((size_t)needed + 1); + if (json == NULL) { + *error_msg = "Failed to allocate array interface JSON"; + goto CLEANUP; + } + + // Build final JSON + int written = snprintf(json, (size_t)needed + 1, + "{\"typestr\":\"%s\",\"shape\":%s," + "\"data\":[%" PRIuPTR ",%s],\"version\":3}", + typestr, shape_json, address, + readonly ? "true" : "false"); + + if (written != needed) { + *error_msg = "Failed to construct array interface JSON"; + goto CLEANUP; + } + + // Success + *json_out = json; + json = NULL; + ok = 1; + +CLEANUP: + if (typestr != NULL) { + enif_free(typestr); + } + + if (shape_json != NULL) { + enif_free(shape_json); + } + + if (json != NULL) { + enif_free(json); + } + + return ok; +} + ERL_NIF_TERM exg_get_int_size(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { ERL_NIF_TERM ret = 0; diff --git a/lib/exgboost.ex b/lib/exgboost.ex index a74ea77..d7a30b4 100644 --- a/lib/exgboost.ex +++ b/lib/exgboost.ex @@ -233,12 +233,15 @@ defmodule EXGBoost do case data do %Nx.Tensor{} = data -> - data_interface = ArrayInterface.from_tensor(data) |> Jason.encode!() + arr = ArrayInterface.from_tensor(data) {shape, preds} = EXGBoost.NIF.booster_predict_from_dense( boostr.ref, - data_interface, + arr.binary, + arr.typestr, + Tuple.to_list(arr.shape), + arr.readonly, Jason.encode!(params), proxy ) @@ -247,16 +250,25 @@ defmodule EXGBoost do 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_arr = ArrayInterface.from_tensor(indptr) + indices_arr = ArrayInterface.from_tensor(indices) + values_arr = ArrayInterface.from_tensor(values) {shape, preds} = EXGBoost.NIF.booster_predict_from_csr( boostr.ref, - indptr_interface, - indices_interface, - values_interface, + indptr_arr.binary, + indptr_arr.typestr, + Tuple.to_list(indptr_arr.shape), + indptr_arr.readonly, + indices_arr.binary, + indices_arr.typestr, + Tuple.to_list(indices_arr.shape), + indices_arr.readonly, + values_arr.binary, + values_arr.typestr, + Tuple.to_list(values_arr.shape), + values_arr.readonly, ncol, Jason.encode!(params), proxy @@ -267,12 +279,15 @@ defmodule EXGBoost do data -> data = Nx.concatenate(data) - data_interface = ArrayInterface.from_tensor(data) |> Jason.encode!() + arr = ArrayInterface.from_tensor(data) {shape, preds} = EXGBoost.NIF.booster_predict_from_dense( boostr.ref, - data_interface, + arr.binary, + arr.typestr, + Tuple.to_list(arr.shape), + arr.readonly, Jason.encode!(params), proxy ) diff --git a/lib/exgboost/array_interface.ex b/lib/exgboost/array_interface.ex index ad2fc8b..05541b7 100644 --- a/lib/exgboost/array_interface.ex +++ b/lib/exgboost/array_interface.ex @@ -1,10 +1,9 @@ defmodule EXGBoost.ArrayInterface do @moduledoc false - alias EXGBoost.Internal @typedoc """ The XGBoost C API uses and is moving towards mainly supporting the use of - JSON-Encoded NumPy ArrayyInterface format to pass data to and from the C API. This struct + JSON-Encoded NumPy ArrayInterface format to pass data to and from the C API. This struct is used to represent the ArrayInterface format. If you wish to use the EXGBoost.NIF library directly, this will be the desired format @@ -17,46 +16,19 @@ defmodule EXGBoost.ArrayInterface do @type t :: %__MODULE__{ typestr: String.t(), shape: tuple(), - address: pos_integer(), readonly: boolean(), - tensor: Nx.Tensor.t(), - binary: binary() + binary: binary() | nil } - @enforce_keys [:typestr, :shape, :address, :readonly] + @enforce_keys [:typestr, :shape, :readonly] defstruct [ :typestr, :shape, - :address, :readonly, - :tensor, :binary, version: 3 ] - defimpl Jason.Encoder do - def encode( - %{ - typestr: typestr, - shape: shape, - address: address, - readonly: readonly, - version: version - }, - opts - ) do - Jason.Encode.map( - %{ - typestr: typestr, - shape: Tuple.to_list(shape), - data: [address, readonly], - version: version - }, - opts - ) - end - end - defimpl Inspect do import Inspect.Algebra @@ -64,7 +36,6 @@ defmodule EXGBoost.ArrayInterface do %{ typestr: typestr, shape: shape, - address: address, readonly: readonly, version: version }, @@ -77,7 +48,7 @@ defmodule EXGBoost.ArrayInterface do %{ typestr: typestr, shape: Tuple.to_list(shape), - data: [address, readonly], + readonly: readonly, version: version }, opts @@ -91,14 +62,38 @@ defmodule EXGBoost.ArrayInterface do def from_map(%{} = interface) do interface |> Enum.reduce([], fn - {"data", [address, readonly]}, acc -> - [{:address, address} | [{:readonly, readonly} | acc]] + {"data", [_address, readonly]}, acc -> + [{:readonly, readonly} | acc] + + {:data, [_address, readonly]}, acc -> + [{:readonly, readonly} | acc] {"shape", shape}, acc -> [{:shape, List.to_tuple(shape)} | acc] - {key, value}, acc -> - [{String.to_existing_atom(key), value} | acc] + {:shape, shape}, acc -> + [{:shape, List.to_tuple(shape)} | acc] + + {"typestr", typestr}, acc -> + [{:typestr, typestr} | acc] + + {:typestr, typestr}, acc -> + [{:typestr, typestr} | acc] + + {"version", version}, acc -> + [{:version, version} | acc] + + {:version, version}, acc -> + [{:version, version} | acc] + + {"binary", binary}, acc -> + [{:binary, binary} | acc] + + {:binary, binary}, acc -> + [{:binary, binary} | acc] + + {_key, _value}, acc -> + acc end) |> then(&struct(__MODULE__, &1)) end @@ -107,9 +102,10 @@ defmodule EXGBoost.ArrayInterface do This function is used to convert Nx.Tensors to the ArrayInterface format. Example: - iex> EXGBoost.from_tensor(Nx.tensor([[1,2,3],[4,5,6]])) + iex> EXGBoost.ArrayInterface.from_tensor(Nx.tensor([[1,2,3],[4,5,6]])) #ArrayInterface< - %{data: [4418559984, true], shape: [2, 3], typestr: " """ @spec from_tensor(Nx.Tensor.t()) :: %__MODULE__{} def from_tensor(%Nx.Tensor{type: t_type} = tensor) do @@ -128,42 +124,39 @@ defmodule EXGBoost.ArrayInterface do end binary = Nx.to_binary(tensor) - tensor_addr = - EXGBoost.NIF.get_binary_address(binary) |> EXGBoost.Internal.unwrap!() %__MODULE__{ typestr: type_char, shape: Nx.shape(tensor), - address: tensor_addr, readonly: true, - tensor: tensor, binary: 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() - <<_endianess::utf8, char_code::binary-size(1), bytes::binary>> = arr_int.typestr + def get_tensor(%__MODULE__{binary: binary, typestr: typestr, shape: shape}) when is_binary(binary) do + # Parse typestr to get Nx type + <<_endian::utf8, char_code::binary-size(1), bytes::binary>> = typestr + bit_width = String.to_integer(bytes) * 8 nx_type = case char_code do - "i" -> {:s, String.to_integer(bytes) * 8} - other -> {String.to_existing_atom(other), String.to_integer(bytes) * 8} + "i" -> {:s, bit_width} + "u" -> {:u, bit_width} + "f" -> {:f, bit_width} + "c" -> {:c, bit_width} end - tensor_bin = - EXGBoost.NIF.get_binary_from_address(arr_int.address, String.to_integer(bytes) * num_items) - |> Internal.unwrap!() - - Nx.from_binary( - tensor_bin, - nx_type - ) - |> Nx.reshape(arr_int.shape) + Nx.from_binary(binary, nx_type) |> Nx.reshape(shape) end - def get_tensor(%__MODULE__{tensor: %Nx.Tensor{} = tensor}) do - tensor + def get_tensor(%__MODULE__{binary: nil}) do + raise ArgumentError, """ + Cannot reconstruct tensor from ArrayInterface without binary data. + + ArrayInterface instances must include the binary data field. + If you're seeing this error, ensure the ArrayInterface was created with + from_tensor/1 or includes binary data from a NIF (like get_quantile_cut). + """ end end diff --git a/lib/exgboost/booster.ex b/lib/exgboost/booster.ex index 1724ee2..adbdac4 100644 --- a/lib/exgboost/booster.ex +++ b/lib/exgboost/booster.ex @@ -294,7 +294,7 @@ defmodule EXGBoost.Booster do EXGBoost.NIF.booster_load_json_config(booster.ref, config) |> Internal.unwrap!() :model -> - EXGBoost.NIF.booster_deserialize_from_buffer(source) |> Internal.unwrap!() + load_model_with_fallback(source) :weights -> EXGBoost.NIF.booster_load_model_from_buffer(source) |> Internal.unwrap!() @@ -303,6 +303,16 @@ defmodule EXGBoost.Booster do struct(booster, ref: booster_ref) end + defp load_model_with_fallback(source) do + case EXGBoost.NIF.booster_load_model_from_buffer(source) do + {:error, _reason} -> + EXGBoost.NIF.booster_deserialize_from_buffer(source) |> Internal.unwrap!() + + result -> + Internal.unwrap!(result) + end + end + @doc """ Get a formatted representation of the Booster's model. diff --git a/lib/exgboost/dmatrix.ex b/lib/exgboost/dmatrix.ex index 141cfc1..6bf3201 100644 --- a/lib/exgboost/dmatrix.ex +++ b/lib/exgboost/dmatrix.ex @@ -127,12 +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!() + arr = ArrayInterface.from_tensor(value) + shape = Tuple.to_list(arr.shape) EXGBoost.NIF.dmatrix_set_info_from_interface( dmat.ref, Atom.to_string(key), - data_interface + arr.binary, + arr.typestr, + shape, + arr.readonly ) end) @@ -164,23 +168,36 @@ defmodule EXGBoost.DMatrix do # config – JSON configuration string. At the moment it should be an empty document, preserved for future use. config = %{} |> Jason.encode!() - {indptr, data} = + {indptr_data, data_data} = EXGBoost.NIF.dmatrix_get_quantile_cut(dmat.ref, config) |> Internal.unwrap!() - indptr = - Jason.decode!(indptr) - |> ArrayInterface.from_map() - |> ArrayInterface.get_tensor() - - data = - Jason.decode!(data) - |> ArrayInterface.from_map() - |> ArrayInterface.get_tensor() + # NIF now returns maps with :binary, :typestr, :shape - no JSON parsing needed + indptr = build_tensor_from_map(indptr_data) + data = build_tensor_from_map(data_data) {indptr, data} end + defp build_tensor_from_map(%{binary: binary, typestr: typestr, shape: shape}) do + # Parse typestr to get Nx type + <<_endian::utf8, char_code::binary-size(1), bytes::binary>> = typestr + bit_width = String.to_integer(bytes) * 8 + + nx_type = + case char_code do + "i" -> {:s, bit_width} + "u" -> {:u, bit_width} + "f" -> {:f, bit_width} + "c" -> {:c, bit_width} + end + + # Convert shape list to tuple + shape_tuple = List.to_tuple(shape) + + Nx.from_binary(binary, nx_type) |> Nx.reshape(shape_tuple) + end + defimpl Inspect do import Inspect.Algebra alias EXGBoost.DMatrix @@ -208,9 +225,16 @@ defmodule EXGBoost.DMatrix do end end + @deprecated "Since 0.9.1 -- Read the file yourself and use `EXGBoost.DMatrix.from_tensor/2` instead" @doc """ Create a DMatrix from a file. + > #### Deprecated {: .warning} + > + > XGBoost deprecated text file input (LIBSVM and CSV) in 3.1 and will remove it in a + > future release. Loading a file through this function emits a deprecation warning from + > XGBoost itself. Parse the file into an `Nx.Tensor` and use `from_tensor/2` instead. + Refer to https://xgboost.readthedocs.io/en/latest/tutorials/external_memory.html#text-file-inputs for proper formatting of the file and the options. @@ -304,9 +328,15 @@ defmodule EXGBoost.DMatrix do config = Enum.into(config_opts, %{}, fn {key, value} -> {Atom.to_string(key), value} end) format = Keyword.fetch!(format_opts, :format) + arr = ArrayInterface.from_tensor(tensor) + shape = Tuple.to_list(arr.shape) + dmat = EXGBoost.NIF.dmatrix_create_from_dense( - Jason.encode!(ArrayInterface.from_tensor(tensor)), + arr.binary, + arr.typestr, + shape, + arr.readonly, Jason.encode!(config) ) |> Internal.unwrap!() @@ -363,11 +393,28 @@ defmodule EXGBoost.DMatrix do raise ArgumentError, "Sparse format must be :csr or :csc" end + indptr_arr = ArrayInterface.from_tensor(indptr) + indices_arr = ArrayInterface.from_tensor(indices) + data_arr = ArrayInterface.from_tensor(data) + + indptr_shape = Tuple.to_list(indptr_arr.shape) + indices_shape = Tuple.to_list(indices_arr.shape) + data_shape = Tuple.to_list(data_arr.shape) + 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)), + indptr_arr.binary, + indptr_arr.typestr, + indptr_shape, + indptr_arr.readonly, + indices_arr.binary, + indices_arr.typestr, + indices_shape, + indices_arr.readonly, + data_arr.binary, + data_arr.typestr, + data_shape, + data_arr.readonly, n, Jason.encode!(config), Atom.to_string(format) diff --git a/lib/exgboost/internal.ex b/lib/exgboost/internal.ex index dd4a3a1..0d24ba0 100644 --- a/lib/exgboost/internal.ex +++ b/lib/exgboost/internal.ex @@ -139,6 +139,17 @@ defmodule EXGBoost.Internal do end def unwrap!({:ok, val}), do: val - def unwrap!({:error, reason}), do: raise(reason) + + def unwrap!({:error, reason}) when is_exception(reason) do + raise(reason) + end + + def unwrap!({:error, reason}) when is_list(reason) do + message = if List.ascii_printable?(reason), do: List.to_string(reason), else: inspect(reason) + raise ArgumentError, message + end + + def unwrap!({:error, reason}) when is_binary(reason), do: raise(ArgumentError, reason) + def unwrap!({:error, reason}), do: raise(ArgumentError, inspect(reason)) def unwrap!(:ok), do: :ok end diff --git a/lib/exgboost/nif.ex b/lib/exgboost/nif.ex index bbed9e0..51f433a 100644 --- a/lib/exgboost/nif.ex +++ b/lib/exgboost/nif.ex @@ -15,7 +15,9 @@ defmodule EXGBoost.NIF do JSON-Encoded Array Interface as defined in the NumPy documentation. https://numpy.org/doc/stable/reference/arrays.interface.html """ - @type array_interface :: String.t() + + @type typestr :: String.t() + @type shape :: [integer()] @type dmatrix_reference :: reference() @type booster_reference :: reference() @type exgboost_return_type(return_type) :: {:ok, return_type} | {:error, String.t()} @@ -100,6 +102,13 @@ defmodule EXGBoost.NIF do def dmatrix_create_from_file(_file_uri, _silent), do: :erlang.nif_error(:not_implemented) + @doc """ + Create a DMatrix from a URI. + + XGBoost deprecated text file input (LIBSVM and CSV) in 3.1 and logs a warning the first + time this is called in a process. Prefer building a DMatrix from an in-memory tensor + with `EXGBoost.DMatrix.from_tensor/2`. + """ def dmatrix_create_from_uri(_config), do: :erlang.nif_error(:not_implemented) @spec dmatrix_create_from_mat(binary, integer(), integer(), float()) :: @@ -120,44 +129,69 @@ defmodule EXGBoost.NIF do do: :erlang.nif_error(:not_implemented) @spec dmatrix_create_from_sparse( - array_interface(), - array_interface(), - array_interface(), + binary(), + typestr(), + shape(), + boolean(), + binary(), + typestr(), + shape(), + boolean(), + binary(), + typestr(), + shape(), + boolean(), integer(), String.t(), String.t() ) :: exgboost_return_type(dmatrix_reference()) @doc """ - Create a DMatrix from a Sparse matrix (CSR / CSC) - - Returns a reference to the DMatrix. - - ## Examples - - iex> EXGBoost.NIF.dmatrix_create_from_csr([0, 2, 3], [0, 2, 2, 0], [1, 2, 3, 4], 2, 2, -1.0) - {:ok, #Reference<>} - - iex> EXGBoost.NIF.dmatrix_create_from_csr([0, 2, 3], [0, 2, 2, 0], [1, 2, 3, 4], 2, 2, -1.0) - {:error #Reference<>} + Create a DMatrix from Sparse matrix Array Interface components (CSR / CSC) + + Arguments for each array (indptr, indices, data): + - binary: The raw tensor data + - typestr: Data type string (e.g., \" Jason.decode!() - |> ExJsonSchema.Schema.resolve() + @schema (fn -> + {:ok, {{_, 200, _}, _, body}} = + :httpc.request(:get, {"https://vega.github.io/schema/vega/v5.json", []}, [], []) + + body |> to_string() |> Jason.decode!() |> ExJsonSchema.Schema.resolve() + end).() @mark_text_doc "Accepts a keyword list of Vega `text` Mark properties. Reference [here](https://vega.github.io/vega/docs/marks/text/) for more details. Accepts either a string (expected to be valid Vega property names) or Elixir-styled atom. Note that keys are snake-cased instead of camel-case (e.g. Vega `fontSize` becomes `font_size`)" @mark_rect_doc "Accepts a keyword list of Vega `rect` Mark properties. Reference [here](https://vega.github.io/vega/docs/marks/rect/) for more details. Accepts either a string (expected to be valid Vega property names) or Elixir-styled atom. Note that keys are snake-cased instead of camel-case (e.g. Vega `fontSize` becomes `font_size`)" @@ -797,15 +801,6 @@ defmodule EXGBoost.Plotting do to ensure that it is valid. Note that this will only validate the specification against the Vega schema, and not against the VegaLite schema. This requires the [`ex_json_schema`] package to be installed. - ## Livebook Integration - - This module also provides a `Kino.Render` implementation for `EXGBoost.Booster` which allows - models to be rendered directly in Livebook. This is done by converting the model into a Vega specification - and then using the `Kino.Render` implementation for Elixir's [`VegaLite`](https://hexdocs.pm/vega_lite/VegaLite.html) API - to render the model. - - . The Vega specification is then passed to [VegaLite](https://hexdocs.pm/vega_lite/readme.html) - ## Plotting Parameters This module exposes a high-level API for customizing the EXGBoost model visualization, but it is also possible to @@ -834,7 +829,7 @@ defmodule EXGBoost.Plotting do @spec get_defaults() :: Keyword.t() def get_defaults(), do: @defaults - @spec get_styles() :: [{atom(), [style(), ...]}, ...] + @spec get_styles() :: EXGBoost.Plotting.Style.t() def get_styles(), do: @styles defp validate_spec(spec) do @@ -1766,9 +1761,3 @@ defmodule EXGBoost.Plotting do IO.inspect(env) end end - -defimpl Kino.Render, for: EXGBoost.Booster do - def to_livebook(booster) do - EXGBoost.Plotting.plot(booster) |> Kino.Render.to_livebook() - end -end diff --git a/lib/exgboost/training/callback.ex b/lib/exgboost/training/callback.ex index b3e6287..ee1d28a 100644 --- a/lib/exgboost/training/callback.ex +++ b/lib/exgboost/training/callback.ex @@ -43,6 +43,14 @@ defmodule EXGBoost.Training.Callback do """ alias EXGBoost.Training.State + + @type t :: %__MODULE__{ + event: event(), + fun: fun(), + name: atom(), + init_state: any() + } + @enforce_keys [:event, :fun] defstruct [:event, :fun, :name, :init_state] @@ -54,7 +62,7 @@ defmodule EXGBoost.Training.Callback do @doc """ Factory for a new callback with an initial state. """ - @spec new(event :: event(), fun :: fun(), name :: atom(), init_state :: any()) :: Callback.t() + @spec new(event :: event(), fun :: fun(), name :: atom(), init_state :: any()) :: t() def new(event, fun, name, init_state \\ %{}) when event in @valid_events and is_function(fun, 1) and is_atom(name) and not is_nil(name) do %__MODULE__{event: event, fun: fun, name: name, init_state: init_state} diff --git a/lib/exgboost/training/state.ex b/lib/exgboost/training/state.ex index de7cd90..4741409 100644 --- a/lib/exgboost/training/state.ex +++ b/lib/exgboost/training/state.ex @@ -1,5 +1,15 @@ defmodule EXGBoost.Training.State do @moduledoc false + + @type t :: %__MODULE__{ + booster: EXGBoost.Booster.t(), + iteration: non_neg_integer(), + max_iteration: -1 | non_neg_integer(), + meta_vars: map(), + metrics: map(), + status: :cont | :halt + } + @enforce_keys [:booster] defstruct [ :booster, diff --git a/mix.exs b/mix.exs index 9a12ab9..c11f62d 100644 --- a/mix.exs +++ b/mix.exs @@ -1,6 +1,7 @@ defmodule EXGBoost.MixProject do use Mix.Project - @version "0.5.2" + + @version "0.9.1" def project do [ @@ -8,34 +9,34 @@ defmodule EXGBoost.MixProject do version: @version, make_precompiler: {:nif, CCPrecompiler}, make_precompiler_url: - "https://github.com/acalejos/exgboost/releases/download/v#{@version}/@{artefact_filename}", + "https://github.com/iperks/exgboost/releases/download/v#{@version}/@{artefact_filename}", make_precompiler_priv_paths: ["libexgboost.*", "lib"], # NIF Versions correspond to OTP Releases # https://github.com/erlang/otp/blob/d3aa6c044c3927f011fb76ac087d5ce0e814954c/erts/emulator/beam/erl_nif.h#L57 make_precompiler_nif_versions: [ - versions: ["2.15", "2.16", "2.17"] + versions: ["2.15", "2.16", "2.17", "2.18"] ], - elixir: "~> 1.14", + elixir: "~> 1.17", start_permanent: Mix.env() == :prod, compilers: [:elixir_make] ++ Mix.compilers(), deps: deps(), name: "EXGBoost", - source_url: "https://github.com/acalejos/exgboost", - homepage_url: "https://github.com/acalejos/exgboost", + source_url: "https://github.com/iperks/exgboost", + homepage_url: "https://github.com/iperks/exgboost", docs: docs(), package: package(), - preferred_cli_env: [ - docs: :docs, - "hex.publish": :docs - ], before_closing_body_tag: &before_closing_body_tag/1, name: "EXGBoost", description: "Elixir bindings for the XGBoost library. `EXGBoost` provides an implementation of XGBoost that works with - [Nx](https://hexdocs.pm/nx/Nx.html) tensors." + [Nx](https://hexdocs.pm/nx/Nx.html) tensors. Maintained fork of acalejos/exgboost." ] end + def cli do + [preferred_envs: [docs: :docs, "hex.publish": :docs]] + end + def application do [ extra_applications: [:logger], @@ -47,26 +48,24 @@ defmodule EXGBoost.MixProject do [ {:elixir_make, "~> 0.4", runtime: false}, {:nimble_options, "~> 1.0"}, - {:nx, "~> 0.7"}, + {:nx, "~> 0.9"}, {:jason, "~> 1.3"}, - {:ex_doc, "~> 0.31.0", only: :docs}, + {:ex_doc, "~> 0.40", only: :docs}, {:cc_precompiler, "~> 0.1.0", runtime: false}, {:exterval, "0.2.0"}, - {:ex_json_schema, "~> 0.11.0"}, - {:httpoison, "~> 2.0", runtime: false}, + {:ex_json_schema, "~> 0.11.4"}, {:vega_lite, "~> 0.1"}, {:vega_lite_convert, "~> 1.0.1"}, - {:kino, "~> 0.11"}, {:scidata, "~> 0.1", only: :dev}, - {:kino_vega_lite, "~> 0.1.9", only: :dev} + {:dialyxir, "~> 1.4", only: [:dev, :test], runtime: false} ] end defp package do [ - maintainers: ["Andres Alejos"], + maintainers: ["Ian Perks"], licenses: ["Apache-2.0"], - links: %{"GitHub" => "https://github.com/acalejos/exgboost"}, + links: %{"GitHub" => "https://github.com/iperks/exgboost"}, files: [ "lib", "mix.exs", diff --git a/mix.lock b/mix.lock index 452af39..87c96c4 100644 --- a/mix.lock +++ b/mix.lock @@ -1,37 +1,27 @@ %{ "castore": {:hex, :castore, "0.1.22", "4127549e411bedd012ca3a308dede574f43819fe9394254ca55ab4895abfa1a2", [:mix], [], "hexpm", "c17576df47eb5aa1ee40cc4134316a99f5cad3e215d5c77b8dd3cfef12a22cac"}, "cc_precompiler": {:hex, :cc_precompiler, "0.1.11", "8c844d0b9fb98a3edea067f94f616b3f6b29b959b6b3bf25fee94ffe34364768", [:mix], [{:elixir_make, "~> 0.7", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "3427232caf0835f94680e5bcf082408a70b48ad68a5f5c0b02a3bea9f3a075b9"}, - "certifi": {:hex, :certifi, "2.15.0", "0e6e882fcdaaa0a5a9f2b3db55b1394dba07e8d6d9bcad08318fb604c6839712", [:rebar3], [], "hexpm", "b147ed22ce71d72eafdad94f055165c1c182f61a2ff49df28bcc71d1d5b94a60"}, - "complex": {:hex, :complex, "0.6.0", "b0130086a7a8c33574d293b2e0e250f4685580418eac52a5658a4bd148f3ccf1", [:mix], [], "hexpm", "0a5fa95580dcaf30fcd60fe1aaf24327c0fe401e98c24d892e172e79498269f9"}, - "decimal": {:hex, :decimal, "2.3.0", "3ad6255aa77b4a3c4f818171b12d237500e63525c2fd056699967a3e7ea20f62", [:mix], [], "hexpm", "a4d66355cb29cb47c3cf30e71329e58361cfcb37c34235ef3bf1d7bf3773aeac"}, - "earmark_parser": {:hex, :earmark_parser, "1.4.44", "f20830dd6b5c77afe2b063777ddbbff09f9759396500cdbe7523efd58d7a339c", [:mix], [], "hexpm", "4778ac752b4701a5599215f7030989c989ffdc4f6df457c5f36938cc2d2a2750"}, - "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"}, - "ex_doc": {:hex, :ex_doc, "0.31.2", "8b06d0a5ac69e1a54df35519c951f1f44a7b7ca9a5bb7a260cd8a174d6322ece", [:mix], [{:earmark_parser, "~> 1.4.39", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.1", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1", [hex: :makeup_erlang, repo: "hexpm", optional: false]}], "hexpm", "317346c14febaba9ca40fd97b5b5919f7751fb85d399cc8e7e8872049f37e0af"}, - "ex_json_schema": {:hex, :ex_json_schema, "0.11.1", "b593f92937a095f66054bb318681397dfe7304e7d2b6b1a7534ea3aa40024f8c", [:mix], [{:decimal, "~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}], "hexpm", "32d651a575a6ce2fd613f140b0fef8dd0acc7cf8e8bcd29a3a1be5c945700dd5"}, + "complex": {:hex, :complex, "0.7.0", "695632ef9487517aa5d57edd1697801079d622414cb2e1a7cf538b1f9a50f205", [:mix], [], "hexpm", "0ee39c0803129f546e7f3f640da8f021c9e659402bf59da6f7f2c4848f068f8d"}, + "decimal": {:hex, :decimal, "3.1.1", "430d87b04011ce6cbd4fd205be758311a81f87d552d40904abd00f015935b1d0", [:mix], [], "hexpm", "c5f25f2ced74a0587d03e6023f595db8e924c9d3922c8c8ffd9edfc4498cf1f6"}, + "dialyxir": {:hex, :dialyxir, "1.4.7", "dda948fcee52962e4b6c5b4b16b2d8fa7d50d8645bbae8b8685c3f9ecb7f5f4d", [:mix], [{:erlex, ">= 0.2.8", [hex: :erlex, repo: "hexpm", optional: false]}], "hexpm", "b34527202e6eb8cee198efec110996c25c5898f43a4094df157f8d28f27d9efe"}, + "earmark_parser": {:hex, :earmark_parser, "1.4.45", "cba8369ab2a1342e419bc2760eec731b17be828941dcf494045d44766227e1d5", [:mix], [], "hexpm", "d3ec045bf122965db20c0bdb420e19ee1415843135327124918473feb4b328e8"}, + "elixir_make": {:hex, :elixir_make, "0.10.0", "16577e2583a79bb79237bbff349619ef5d80afffc07eac6e4faf0d00e2ddaf7d", [:mix], [], "hexpm", "dc1f09fb7fa68866b886abd5f0f3c83553b1a19a52359a899e92af1bb3b31982"}, + "erlex": {:hex, :erlex, "0.2.9", "7debbbaa9f4f368b8cd648983e0f1d7963028508e9c59e9d4ed504e94ef52a55", [:mix], [], "hexpm", "8cfffc0ec7159e6d73de2ab28a588064de80f88b2798d5cbe4482cbbc200178b"}, + "ex_doc": {:hex, :ex_doc, "0.40.3", "4a972ffe64bc07dc605af487e98fc19b72a4185f55ca031b94c0552d6071c1d9", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "2756e357742fecd9749b489b85d67c9ce99c465f2e75728d9e6dc8d704b973de"}, + "ex_json_schema": {:hex, :ex_json_schema, "0.11.5", "ea45f3238be135949dbbbcc9e8eb4682d6e561b8a66374e75d85a2e1d2bd4107", [:mix], [{:decimal, "~> 3.0", [hex: :decimal, repo: "hexpm", optional: false]}], "hexpm", "61ed2a8f07bd115e7ab6d45c147642a8c73b962bc419fadbb248046b9d3d0f20"}, "exterval": {:hex, :exterval, "0.2.0", "cb2ce1dcaa08f7827d5eab25e0268e786b6b05b9b6d27f79d1d77f7b9f6dbd5d", [:mix], [], "hexpm", "30b48089fcc9a648d4bfcd17d51ab799b846c34db0fbf1fd6eb8a5526587510d"}, - "fss": {:hex, :fss, "0.1.1", "9db2344dbbb5d555ce442ac7c2f82dd975b605b50d169314a20f08ed21e08642", [:mix], [], "hexpm", "78ad5955c7919c3764065b21144913df7515d52e228c09427a004afe9c1a16b0"}, - "hackney": {:hex, :hackney, "1.25.0", "390e9b83f31e5b325b9f43b76e1a785cbdb69b5b6cd4e079aa67835ded046867", [:rebar3], [{:certifi, "~> 2.15.0", [hex: :certifi, repo: "hexpm", optional: false]}, {:idna, "~> 6.1.0", [hex: :idna, repo: "hexpm", optional: false]}, {:metrics, "~> 1.0.0", [hex: :metrics, repo: "hexpm", optional: false]}, {:mimerl, "~> 1.4", [hex: :mimerl, repo: "hexpm", optional: false]}, {:parse_trans, "3.4.1", [hex: :parse_trans, repo: "hexpm", optional: false]}, {:ssl_verify_fun, "~> 1.1.0", [hex: :ssl_verify_fun, repo: "hexpm", optional: false]}, {:unicode_util_compat, "~> 0.7.1", [hex: :unicode_util_compat, repo: "hexpm", optional: false]}], "hexpm", "7209bfd75fd1f42467211ff8f59ea74d6f2a9e81cbcee95a56711ee79fd6b1d4"}, - "httpoison": {:hex, :httpoison, "2.2.3", "a599d4b34004cc60678999445da53b5e653630651d4da3d14675fedc9dd34bd6", [:mix], [{:hackney, "~> 1.21", [hex: :hackney, repo: "hexpm", optional: false]}], "hexpm", "fa0f2e3646d3762fdc73edb532104c8619c7636a6997d20af4003da6cfc53e53"}, - "idna": {:hex, :idna, "6.1.1", "8a63070e9f7d0c62eb9d9fcb360a7de382448200fbbd1b106cc96d3d8099df8d", [:rebar3], [{:unicode_util_compat, "~> 0.7.0", [hex: :unicode_util_compat, repo: "hexpm", optional: false]}], "hexpm", "92376eb7894412ed19ac475e4a86f7b413c1b9fbb5bd16dccd57934157944cea"}, - "jason": {:hex, :jason, "1.4.4", "b9226785a9aa77b6857ca22832cffa5d5011a667207eb2a0ad56adb5db443b8a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "c5eb0cab91f094599f94d55bc63409236a8ec69a21a67814529e8d5f6cc90b3b"}, - "kino": {:hex, :kino, "0.17.0", "72f1a2bf691db7b8352bae86b3951fdf9b23619b5d8586cb7cd1e9c2edc8ff9b", [:mix], [{:fss, "~> 0.1.0", [hex: :fss, repo: "hexpm", optional: false]}, {:nx, "~> 0.1", [hex: :nx, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}, {:table, "~> 0.1.2", [hex: :table, repo: "hexpm", optional: false]}], "hexpm", "e1ec49a2ebbf622c1675f96b427c565ce02df6725e8f2e8d4a743c8e791bd090"}, - "kino_vega_lite": {:hex, :kino_vega_lite, "0.1.13", "03c00405987a2202e4b8014ee55eb7f5727691b3f13d76a3764f6eeccef45322", [:mix], [{:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: false]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: false]}, {:vega_lite, "~> 0.1.8", [hex: :vega_lite, repo: "hexpm", optional: false]}], "hexpm", "00c72bc270e7b9d3c339f726cdab0012fd3f2fc75e36c7548e0f250fe420fa10"}, - "makeup": {:hex, :makeup, "1.2.1", "e90ac1c65589ef354378def3ba19d401e739ee7ee06fb47f94c687016e3713d1", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "d36484867b0bae0fea568d10131197a4c2e47056a6fbe84922bf6ba71c8d17ce"}, - "makeup_elixir": {:hex, :makeup_elixir, "0.16.2", "627e84b8e8bf22e60a2579dad15067c755531fea049ae26ef1020cad58fe9578", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "41193978704763f6bbe6cc2758b84909e62984c7752b3784bd3c218bb341706b"}, - "makeup_erlang": {:hex, :makeup_erlang, "0.1.5", "e0ff5a7c708dda34311f7522a8758e23bfcd7d8d8068dc312b5eb41c6fd76eba", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "94d2e986428585a21516d7d7149781480013c56e30c6a233534bedf38867a59a"}, - "metrics": {:hex, :metrics, "1.0.1", "25f094dea2cda98213cecc3aeff09e940299d950904393b2a29d191c346a8486", [:rebar3], [], "hexpm", "69b09adddc4f74a40716ae54d140f93beb0fb8978d8636eaded0c31b6f099f16"}, - "mimerl": {:hex, :mimerl, "1.4.0", "3882a5ca67fbbe7117ba8947f27643557adec38fa2307490c4c4207624cb213b", [:rebar3], [], "hexpm", "13af15f9f68c65884ecca3a3891d50a7b57d82152792f3e19d88650aa126b144"}, + "jason": {:hex, :jason, "1.4.5", "2e3a008590b0b8d7388c20293e9dcc9cf3e5d642fd2a114e4cbbb52e595d940a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "b0c823996102bcd0239b3c2444eb00409b72f6a140c1950bc8b457d836b30684"}, + "makeup": {:hex, :makeup, "1.2.2", "882d46dc0905e9ff7abf2aab61a7e6b3dcc555533977d8a23b06019e6c89ac94", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "9a1a24e5b343b8ae16abea0822c10a6f75da27af7fa802ada5251f7579bfccfa"}, + "makeup_elixir": {:hex, :makeup_elixir, "1.0.1", "e928a4f984e795e41e3abd27bfc09f51db16ab8ba1aebdba2b3a575437efafc2", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "7284900d412a3e5cfd97fdaed4f5ed389b8f2b4cb49efc0eb3bd10e2febf9507"}, + "makeup_erlang": {:hex, :makeup_erlang, "1.1.0", "835f7e60792e08824cda445639555d7bf1bbbddb1b60b306e33cb6f6db24dc74", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "1cd6780fb1dd1a03979abaed0fe82712b0625118fd5257d3ebbf73f960c73c3c"}, "nimble_csv": {:hex, :nimble_csv, "1.3.0", "b7f998dc62b222bce9596e46f028c7a5af04cb5dde6df2ea197c583227c54971", [:mix], [], "hexpm", "41ccdc18f7c8f8bb06e84164fc51635321e80d5a3b450761c4997d620925d619"}, "nimble_options": {:hex, :nimble_options, "1.1.1", "e3a492d54d85fc3fd7c5baf411d9d2852922f66e69476317787a7b2bb000a61b", [:mix], [], "hexpm", "821b2470ca9442c4b6984882fe9bb0389371b8ddec4d45a9504f00a66f650b44"}, "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"}, - "nx": {:hex, :nx, "0.10.0", "128e4a094cb790f663e20e1334b127c1f2a4df54edfb8b13c22757ec33133b4f", [:mix], [{:complex, "~> 0.6", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "3db8892c124aeee091df0e6fbf8e5bf1b81f502eb0d4f5ba63e6378ebcae7da4"}, - "parse_trans": {:hex, :parse_trans, "3.4.1", "6e6aa8167cb44cc8f39441d05193be6e6f4e7c2946cb2759f015f8c56b76e5ff", [:rebar3], [], "hexpm", "620a406ce75dada827b82e453c19cf06776be266f5a67cff34e1ef2cbb60e49a"}, - "rustler_precompiled": {:hex, :rustler_precompiled, "0.8.3", "4e741024b0b097fe783add06e53ae9a6f23ddc78df1010f215df0c02915ef5a8", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "c23f5f33cb6608542de4d04faf0f0291458c352a4648e4d28d17ee1098cddcc4"}, + "nx": {:hex, :nx, "0.12.1", "6e9fee43a77646d04faad2ba4e449b9e270c6b23e413f203cb4d81e71c0a617f", [:mix], [{:complex, "~> 0.7", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "ee80f6ae898f68bbfe7f30216b06aab10096231ec99ded1208a827256b23d0fb"}, + "rustler_precompiled": {:hex, :rustler_precompiled, "0.9.0", "3a052eda09f3d2436364645cc1f13279cf95db310eb0c17b0d8f25484b233aa0", [:mix], [{:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "471d97315bd3bf7b64623418b3693eedd8e47de3d1cb79a0ac8f9da7d770d94c"}, "scidata": {:hex, :scidata, "0.1.11", "fe3358bac7d740374b4f2a7eff6a1cb02e5ee7f87f7cdb1e8648ad93c533165f", [:mix], [{:castore, "~> 0.1", [hex: :castore, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.1", [hex: :nimble_csv, repo: "hexpm", optional: false]}, {:stb_image, "~> 0.4", [hex: :stb_image, repo: "hexpm", optional: true]}], "hexpm", "90873337a9d5fe880d640517efa93d3c07e46c8ba436de44117f581800549f93"}, - "ssl_verify_fun": {:hex, :ssl_verify_fun, "1.1.7", "354c321cf377240c7b8716899e182ce4890c5938111a1296add3ec74cf1715df", [:make, :mix, :rebar3], [], "hexpm", "fe4c190e8f37401d30167c8c405eda19469f34577987c76dde613e838bbc67f8"}, "table": {:hex, :table, "0.1.2", "87ad1125f5b70c5dea0307aa633194083eb5182ec537efc94e96af08937e14a8", [:mix], [], "hexpm", "7e99bc7efef806315c7e65640724bf165c3061cdc5d854060f74468367065029"}, - "telemetry": {:hex, :telemetry, "1.3.0", "fedebbae410d715cf8e7062c96a1ef32ec22e764197f70cda73d82778d61e7a2", [:rebar3], [], "hexpm", "7015fc8919dbe63764f4b4b87a95b7c0996bd539e0d499be6ec9d7f3875b79e6"}, - "unicode_util_compat": {:hex, :unicode_util_compat, "0.7.1", "a48703a25c170eedadca83b11e88985af08d35f37c6f664d6dcfb106a97782fc", [:rebar3], [], "hexpm", "b3a917854ce3ae233619744ad1e0102e05673136776fb2fa76234f3e03b23642"}, + "telemetry": {:hex, :telemetry, "1.4.2", "a0cb522801dffb1c49fe6e30561badffc7b6d0e180db1300df759faa22062855", [:rebar3], [], "hexpm", "928f6495066506077862c0d1646609eed891a4326bee3126ba54b60af61febb1"}, "vega_lite": {:hex, :vega_lite, "0.1.11", "2b261d21618f6fa9f63bb4542f0262982d2e40aea3f83e935788fe172902b3c2", [:mix], [{:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: false]}], "hexpm", "d18c3f11369c14bdf36ab53010c06bf5505c221cbcb32faac7420cf6926b3c50"}, "vega_lite_convert": {:hex, :vega_lite_convert, "1.0.1", "1cc0309998c10bce9d944ae631938a433e9cad4ccf7344f9a192d9ddcab1bd93", [:mix], [{:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: false]}, {:rustler, ">= 0.0.0", [hex: :rustler, repo: "hexpm", optional: true]}, {:rustler_precompiled, "~> 0.4", [hex: :rustler_precompiled, repo: "hexpm", optional: false]}, {:vega_lite, ">= 0.0.0", [hex: :vega_lite, repo: "hexpm", optional: false]}], "hexpm", "1e12b4ef3943510d3597e4876a10c6ab333b8ee71204b5e571605d80aa3cbd76"}, } diff --git a/notebooks/plotting.livemd b/notebooks/plotting.livemd index c5c139d..a41ade6 100644 --- a/notebooks/plotting.livemd +++ b/notebooks/plotting.livemd @@ -7,6 +7,13 @@ Mix.install([ {:kino_vega_lite, "~> 0.1"} ]) +# Implement the Kino renderer for later +defimpl Kino.Render, for: EXGBoost.Booster do + def to_livebook(booster) do + EXGBoost.Plotting.plot(booster) |> Kino.Render.to_livebook() + end +end + # This assumed you launch this livebook from its location in the exgboost/notebooks folder ``` diff --git a/scripts/check_xgboost_c_api.sh b/scripts/check_xgboost_c_api.sh new file mode 100755 index 0000000..6ed3fa5 --- /dev/null +++ b/scripts/check_xgboost_c_api.sh @@ -0,0 +1,186 @@ +#!/usr/bin/env bash +set -euo pipefail + +extract_signatures() { + local header="$1" + + awk ' + /XG[A-Za-z0-9_]+[[:space:]]*\(/ { + in_decl=1 + decl=$0 + if ($0 ~ /;/) { + print decl + in_decl=0 + decl="" + } + next + } + in_decl { + decl=decl " " $0 + if ($0 ~ /;/) { + print decl + in_decl=0 + decl="" + } + } + ' "$header" \ + | sed -E 's/[[:space:]]+/ /g; s/^[[:space:]]+//; s/[[:space:]]+$//' \ + | awk ' + { + if (match($0, /XG[A-Za-z0-9_]+[[:space:]]*\(/)) { + name = substr($0, RSTART, RLENGTH) + gsub(/[[:space:]]*\($/, "", name) + print name "\t" $0 + } + } + ' +} + +compare_mode() { + if [[ $# -ne 2 ]]; then + echo "Usage: $0 --compare " + exit 2 + fi + + local old_header="$1/xgboost/c_api.h" + local new_header="$2/xgboost/c_api.h" + + if [[ ! -f "$old_header" ]]; then + echo "XGBoost C API header not found: $old_header" + exit 2 + fi + + if [[ ! -f "$new_header" ]]; then + echo "XGBoost C API header not found: $new_header" + exit 2 + fi + + old_tmp=$(mktemp) + new_tmp=$(mktemp) + trap 'rm -f "${old_tmp:-}" "${new_tmp:-}"' EXIT + + extract_signatures "$old_header" | sort -u > "$old_tmp" + extract_signatures "$new_header" | sort -u > "$new_tmp" + + declare -A old_map=() + declare -A new_map=() + + while IFS=$'\t' read -r name sig; do + [[ -n "${name:-}" ]] && old_map["$name"]="$sig" + done < "$old_tmp" + + while IFS=$'\t' read -r name sig; do + [[ -n "${name:-}" ]] && new_map["$name"]="$sig" + done < "$new_tmp" + + removed=() + added=() + changed=() + + for name in "${!old_map[@]}"; do + if [[ -z "${new_map[$name]+x}" ]]; then + removed+=("$name") + elif [[ "${old_map[$name]}" != "${new_map[$name]}" ]]; then + changed+=("$name") + fi + done + + for name in "${!new_map[@]}"; do + if [[ -z "${old_map[$name]+x}" ]]; then + added+=("$name") + fi + done + + if (( ${#removed[@]} == 0 && ${#changed[@]} == 0 && ${#added[@]} == 0 )); then + echo "No C API declaration differences found between versions." + return 0 + fi + + if (( ${#removed[@]} > 0 )); then + printf 'Removed symbols (%d):\n' "${#removed[@]}" + printf ' - %s\n' "${removed[@]}" + fi + + if (( ${#changed[@]} > 0 )); then + printf 'Changed signatures (%d):\n' "${#changed[@]}" + for name in "${changed[@]}"; do + printf ' - %s\n' "$name" + printf ' old: %s\n' "${old_map[$name]}" + printf ' new: %s\n' "${new_map[$name]}" + done + fi + + if (( ${#added[@]} > 0 )); then + printf 'Added symbols (%d):\n' "${#added[@]}" + printf ' - %s\n' "${added[@]}" + fi + + if (( ${#removed[@]} > 0 || ${#changed[@]} > 0 )); then + return 1 + fi + + return 0 +} + +if [[ "${1:-}" == "--compare" ]]; then + shift + compare_mode "$@" + exit $? +fi + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 [xgboost-shared-lib-path]" + echo " $0 --compare " + exit 2 +fi + +include_dir="$1" +shared_lib="${2:-}" +header="${include_dir}/xgboost/c_api.h" + +if [[ ! -f "$header" ]]; then + echo "XGBoost C API header not found: $header" + exit 2 +fi + +mapfile -t used_symbols < <( + grep -RhoE '\bXG[A-Za-z0-9_]+[[:space:]]*\(' c/exgboost/src/*.c \ + | sed -E 's/[[:space:]]*\($//' \ + | sort -u +) + +missing_in_header=() +for symbol in "${used_symbols[@]}"; do + if ! grep -qE "\\b${symbol}[[:space:]]*\\(" "$header"; then + missing_in_header+=("$symbol") + fi +done + +if (( ${#missing_in_header[@]} > 0 )); then + echo "Symbols used by exgboost NIF but missing from XGBoost C API header:" + printf ' - %s\n' "${missing_in_header[@]}" + exit 1 +fi + +if [[ -n "$shared_lib" && -f "$shared_lib" && "$(command -v nm || true)" != "" ]]; then + if [[ "$(uname -s)" == "Darwin" ]]; then + mapfile -t exported_symbols < <(nm -gU "$shared_lib" | awk '{print $3}' | sed 's/^_//' | sort -u) + else + mapfile -t exported_symbols < <(nm -D --defined-only "$shared_lib" | awk '{print $3}' | sort -u) + fi + + missing_in_lib=() + for symbol in "${used_symbols[@]}"; do + if ! printf '%s\n' "${exported_symbols[@]}" | grep -qx "$symbol"; then + missing_in_lib+=("$symbol") + fi + done + + if (( ${#missing_in_lib[@]} > 0 )); then + echo "Symbols declared in header but not exported by shared library:" + printf ' - %s\n' "${missing_in_lib[@]}" + exit 1 + fi +fi + +echo "XGBoost C API compatibility check passed (${#used_symbols[@]} symbols)." diff --git a/test/exgboost_test.exs b/test/exgboost_test.exs index fcdfec0..3bdd412 100644 --- a/test/exgboost_test.exs +++ b/test/exgboost_test.exs @@ -80,7 +80,7 @@ defmodule EXGBoostTest do assert Booster.get_boosted_rounds(booster) == num_boost_round end - test "predict", context do + test "predict matches inplace_predict", context do nrows = :rand.uniform(10) ncols = :rand.uniform(10) {x, new_key} = Nx.Random.normal(context.key, 0, 1, shape: {nrows, ncols}) @@ -95,6 +95,59 @@ defmodule EXGBoostTest do assert inplace_preds_no_proxy.shape == y.shape end + test "predict is stable across repeated calls", context do + nrows = 12 + ncols = 5 + {x, new_key} = Nx.Random.normal(context.key, 0, 1, shape: {nrows, ncols}) + {y, _new_key} = Nx.Random.normal(new_key, 0, 1, shape: {nrows}) + + booster = + EXGBoost.train(x, y, + num_boost_rounds: 25, + tree_method: :hist, + eval_metric: :rmse + ) + + first_preds = EXGBoost.predict(booster, x) + second_preds = EXGBoost.predict(booster, x) + third_preds = EXGBoost.predict(booster, x) + + assert Nx.all_close(first_preds, second_preds) + assert Nx.all_close(first_preds, third_preds) + assert Nx.all_close(second_preds, third_preds) + end + + test "predict first call matches subsequent confidence values", context do + nrows = 30 + ncols = 5 + {x, new_key} = Nx.Random.normal(context.key, 0, 1, shape: {nrows, ncols}) + {y, _new_key} = Nx.Random.normal(new_key, 0, 1, shape: {nrows}) + + booster = + EXGBoost.train(x, y, + num_boost_rounds: 30, + tree_method: :hist, + eval_metric: :rmse + ) + + sample = Nx.slice_along_axis(x, 0, 1, axis: 0) + + [baseline_confidence] = + EXGBoost.predict(booster, sample) + |> Nx.to_flat_list() + + for i <- 1..10 do + [confidence] = + EXGBoost.predict(booster, sample) + |> Nx.to_flat_list() + + assert_in_delta confidence, + baseline_confidence, + 1.0e-9, + "Expected confidence to match baseline #{baseline_confidence}, got #{confidence} for iteration #{i}" + end + end + test "predict with container", context do nrows = :rand.uniform(10) ncols = :rand.uniform(10) @@ -166,9 +219,6 @@ defmodule EXGBoostTest do ) end) end - - refute is_nil(booster.best_iteration) - refute is_nil(booster.best_score) end test "eval with multiple metrics", context do @@ -280,6 +330,24 @@ defmodule EXGBoostTest do assert is_struct(bst, EXGBoost.Booster) end + test "load_model accepts model artifacts produced by dump_weights", context do + nrows = :rand.uniform(10) + ncols = :rand.uniform(10) + {x, _new_key} = Nx.Random.normal(context[:key], 0, 1, shape: {nrows, ncols}) + {y, _new_key} = Nx.Random.normal(context[:key], 0, 1, shape: {nrows}) + + booster = + EXGBoost.train(x, y, + num_boost_rounds: 10, + tree_method: :hist, + eval_metric: :rmse + ) + + buffer = EXGBoost.dump_weights(booster) + bst = EXGBoost.load_model(buffer) + assert is_struct(bst, EXGBoost.Booster) + end + test "serialize and deserialize weights to and from buffer", context do nrows = :rand.uniform(10) ncols = :rand.uniform(10) @@ -323,12 +391,40 @@ defmodule EXGBoostTest do test "array interface get tensor" do tensor = Nx.tensor([[1, 2, 3], [4, 5, 6]]) array_interface = EXGBoost.ArrayInterface.from_tensor(tensor) - # Set this to nil so we can test the get_tensor reconstruction - array_interface = struct(array_interface, tensor: nil) + # Test that get_tensor returns the cached tensor field assert EXGBoost.ArrayInterface.get_tensor(array_interface) == tensor end + test "array interface from_map ignores optional keys" do + arr_int = + EXGBoost.ArrayInterface.from_map(%{ + "typestr" => " [2, 2], + "data" => [123, true], # Address is ignored, only readonly is extracted + "version" => 3, + "strides" => nil, + "descr" => [["", " + EXGBoost.ArrayInterface.get_tensor(%EXGBoost.ArrayInterface{ + typestr: " Enum.map(fn [_, name, arity] -> {String.to_atom(name), String.to_integer(arity)} end) + |> MapSet.new() + + elixir_exports = + EXGBoost.NIF.__info__(:functions) + |> Enum.reject(fn {name, _arity} -> name in [:module_info, :__info__, :on_load] end) + |> MapSet.new() + + assert c_exports == elixir_exports, + "NIF API mismatch. Missing in C: #{inspect(MapSet.difference(elixir_exports, c_exports))}. " <> + "Missing in Elixir: #{inspect(MapSet.difference(c_exports, elixir_exports))}" + end +end diff --git a/test/nif_test.exs b/test/nif_test.exs index 383418d..78ebc61 100644 --- a/test/nif_test.exs +++ b/test/nif_test.exs @@ -1,5 +1,6 @@ defmodule NifTest do use ExUnit.Case, async: true + import EXGBoost.Internal import EXGBoost.ArrayInterface, only: [from_tensor: 1] @@ -7,6 +8,12 @@ defmodule NifTest do assert EXGBoost.NIF.xgboost_version() |> unwrap!() != :error end + test "unwrap! raises string error for charlist reasons" do + assert_raise ArgumentError, "boom", fn -> + unwrap!({:error, ~c"boom"}) + end + end + test "build_info" do assert EXGBoost.NIF.xgboost_build_info() |> unwrap!() != :error end @@ -28,7 +35,6 @@ defmodule NifTest do end test "dmatrix_create_from_sparse" do - config = Jason.encode!(%{"missing" => 0.0}) indptr = Nx.tensor([0, 22]) ncols = 127 @@ -84,23 +90,45 @@ defmodule NifTest do 1.0 ]) + indptr_arr = from_tensor(indptr) + indices_arr = from_tensor(indices) + data_arr = 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!(), + indptr_arr.binary, + indptr_arr.typestr, + Tuple.to_list(indptr_arr.shape), + indptr_arr.readonly, + indices_arr.binary, + indices_arr.typestr, + Tuple.to_list(indices_arr.shape), + indices_arr.readonly, + data_arr.binary, + data_arr.typestr, + Tuple.to_list(data_arr.shape), + data_arr.readonly, ncols, - config, + config(), "csr" ) |> unwrap!() != :error assert EXGBoost.NIF.dmatrix_create_from_sparse( - from_tensor(indptr) |> Jason.encode!(), - from_tensor(indices) |> Jason.encode!(), - from_tensor(data) |> Jason.encode!(), + indptr_arr.binary, + indptr_arr.typestr, + Tuple.to_list(indptr_arr.shape), + indptr_arr.readonly, + indices_arr.binary, + indices_arr.typestr, + Tuple.to_list(indices_arr.shape), + indices_arr.readonly, + data_arr.binary, + data_arr.typestr, + Tuple.to_list(data_arr.shape), + data_arr.readonly, ncols, - config, + config(), "csc" ) |> unwrap!() != @@ -108,36 +136,54 @@ 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!(), + indptr_arr.binary, + indptr_arr.typestr, + Tuple.to_list(indptr_arr.shape), + indptr_arr.readonly, + indices_arr.binary, + indices_arr.typestr, + Tuple.to_list(indices_arr.shape), + indices_arr.readonly, + data_arr.binary, + data_arr.typestr, + Tuple.to_list(data_arr.shape), + data_arr.readonly, ncols, - config, + config(), "csa" ) assert status == :error end - test "test_dmatrix_create_from_dense" do + test "dmatrix_create_from_dense" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) - - assert EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) - |> unwrap!() != - :error + arr = from_tensor(mat) + shape = Tuple.to_list(arr.shape) + + assert EXGBoost.NIF.dmatrix_create_from_dense( + arr.binary, + arr.typestr, + shape, + arr.readonly, + config() + ) + |> unwrap!() != :error end - test "test_dmatrix_set_str_feature_info" do + test "dmatrix_set_str_feature_info" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + arr = from_tensor(mat) + shape = Tuple.to_list(arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + arr.binary, + arr.typestr, + shape, + arr.readonly, + config() + ) |> unwrap!() assert EXGBoost.NIF.dmatrix_set_str_feature_info(dmat, ~c"feature_name", [ @@ -147,14 +193,19 @@ defmodule NifTest do ]) == :ok end - test "test_dmatrix_get_str_feature_info" do + test "dmatrix_get_str_feature_info" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() EXGBoost.NIF.dmatrix_set_str_feature_info(dmat, ~c"feature_name", [ @@ -168,12 +219,17 @@ defmodule NifTest do test "dmatrix_num_row" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() assert EXGBoost.NIF.dmatrix_num_row(dmat) |> unwrap! == 2 @@ -181,12 +237,17 @@ defmodule NifTest do test "dmatrix_num_col" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() assert EXGBoost.NIF.dmatrix_num_col(dmat) |> unwrap! == 3 @@ -194,12 +255,17 @@ defmodule NifTest do test "dmatrix_num_non_missing" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() assert EXGBoost.NIF.dmatrix_num_non_missing(dmat) |> unwrap! == 6 @@ -207,45 +273,70 @@ defmodule NifTest do test "dmatrix_set_info_from_interface" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() + mat_arr = from_tensor(mat) labels = Nx.tensor([1.0, 0.0]) - - config = Jason.encode!(%{"missing" => -1.0}) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() - label_interface = from_tensor(labels) |> Jason.encode!() + label_arr = from_tensor(labels) + label_shape = Tuple.to_list(label_arr.shape) assert EXGBoost.NIF.dmatrix_set_info_from_interface( dmat, ~c"label", - label_interface + label_arr.binary, + label_arr.typestr, + label_shape, + label_arr.readonly ) == :ok assert EXGBoost.NIF.dmatrix_set_info_from_interface( dmat, ~c"unsupported", - label_interface + label_arr.binary, + label_arr.typestr, + label_shape, + label_arr.readonly ) != :ok end test "dmatrix_save_binary" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() + mat_arr = from_tensor(mat) labels = Nx.tensor([1.0, 0.0]) - - config = Jason.encode!(%{"missing" => -1.0}) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() - interface = from_tensor(labels) |> Jason.encode!() + label_arr = from_tensor(labels) + label_shape = Tuple.to_list(label_arr.shape) - EXGBoost.NIF.dmatrix_set_info_from_interface(dmat, ~c"label", interface) + EXGBoost.NIF.dmatrix_set_info_from_interface( + dmat, + ~c"label", + label_arr.binary, + label_arr.typestr, + label_shape, + label_arr.readonly + ) path = Path.join(System.tmp_dir!(), "test.buffer") |> String.to_charlist() assert EXGBoost.NIF.dmatrix_save_binary(dmat, path, 1) == :ok @@ -253,17 +344,31 @@ defmodule NifTest do test "dmatrix_get_float_info" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() + mat_arr = from_tensor(mat) weights = Nx.tensor([1.0, 0.0]) - - config = Jason.encode!(%{"missing" => -1.0}) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() - interface = from_tensor(weights) |> Jason.encode!() - EXGBoost.NIF.dmatrix_set_info_from_interface(dmat, ~c"feature_weights", interface) + weights_arr = from_tensor(weights) + weights_shape = Tuple.to_list(weights_arr.shape) + + EXGBoost.NIF.dmatrix_set_info_from_interface( + dmat, + ~c"feature_weights", + weights_arr.binary, + weights_arr.typestr, + weights_shape, + weights_arr.readonly + ) assert EXGBoost.NIF.dmatrix_get_float_info(dmat, ~c"feature_weights") |> unwrap!() == Nx.to_list(weights) @@ -271,12 +376,17 @@ defmodule NifTest do test "dmatrix_get_data_as_csr" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() assert EXGBoost.NIF.dmatrix_get_data_as_csr(dmat, Jason.encode!(%{})) |> unwrap!() != :error @@ -284,12 +394,17 @@ defmodule NifTest do test "dmatrix_slice" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() # We do this because the C API uses non fixed-width types so we need to know the size they're expecting from int @@ -323,17 +438,30 @@ defmodule NifTest do test "booster_create" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) mat2 = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - array_interface2 = from_tensor(mat2) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat2_arr = from_tensor(mat2) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() + mat2_shape = Tuple.to_list(mat2_arr.shape) + dmat2 = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface2, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat2_arr.binary, + mat2_arr.typestr, + mat2_shape, + mat2_arr.readonly, + config() + ) |> unwrap!() assert EXGBoost.NIF.booster_create([dmat]) |> unwrap!() != :error @@ -343,26 +471,36 @@ defmodule NifTest do test "booster_get_num_feature" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() booster = EXGBoost.NIF.booster_create([dmat]) |> unwrap!() assert EXGBoost.NIF.booster_get_num_feature(booster) |> unwrap!() == 3 end - test "test_booster_set_str_feature_info" do + test "booster_set_str_feature_info" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() booster = EXGBoost.NIF.booster_create([dmat]) |> unwrap!() @@ -374,14 +512,19 @@ defmodule NifTest do ]) == :ok end - test "test_booster_get_str_feature_info" do + test "booster_get_str_feature_info" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() booster = EXGBoost.NIF.booster_create([dmat]) |> unwrap!() @@ -395,15 +538,20 @@ defmodule NifTest do assert EXGBoost.NIF.booster_get_str_feature_info(booster, ~c"feature_name") |> unwrap!() end - test "test_boster_feature_score" do + test "booster_feature_score" do # TODO: Make more robust test. This will just return an empty list mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() config = Jason.encode!(%{"importance_type" => "weight"}) @@ -414,12 +562,17 @@ defmodule NifTest do test "save model" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() json_file = Path.join(System.tmp_dir!(), "model.json") |> String.to_charlist() @@ -435,12 +588,17 @@ defmodule NifTest do test "load model" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() json_file = Path.join(System.tmp_dir!(), "model.json") |> String.to_charlist() @@ -456,12 +614,17 @@ defmodule NifTest do test "booster serialize" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() booster = EXGBoost.NIF.booster_create([dmat]) |> unwrap!() @@ -470,12 +633,17 @@ defmodule NifTest do test "booster deserialize" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() booster = EXGBoost.NIF.booster_create([dmat]) |> unwrap!() @@ -486,12 +654,17 @@ defmodule NifTest do test "save booster config" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() booster = EXGBoost.NIF.booster_create([dmat]) |> unwrap!() @@ -500,16 +673,163 @@ defmodule NifTest do test "load booster config" do mat = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - array_interface = from_tensor(mat) |> Jason.encode!() - - config = Jason.encode!(%{"missing" => -1.0}) + mat_arr = from_tensor(mat) + mat_shape = Tuple.to_list(mat_arr.shape) dmat = - EXGBoost.NIF.dmatrix_create_from_dense(array_interface, config) + EXGBoost.NIF.dmatrix_create_from_dense( + mat_arr.binary, + mat_arr.typestr, + mat_shape, + mat_arr.readonly, + config() + ) |> unwrap!() booster = EXGBoost.NIF.booster_create([dmat]) |> unwrap!() buf = EXGBoost.NIF.booster_save_json_config(booster) |> unwrap!() assert EXGBoost.NIF.booster_load_json_config(booster, buf) |> unwrap!() != :error end + + describe "Array Interface safety validations" do + test "rejects invalid typestr format" do + binary = <<1, 2, 3, 4, 5, 6, 7, 8>> + + result = EXGBoost.NIF.dmatrix_create_from_dense( + binary, + "invalid_typestr", # Invalid format - should be like "> + + result = EXGBoost.NIF.dmatrix_create_from_dense( + binary, + "> # Only 8 bytes + + result = EXGBoost.NIF.dmatrix_create_from_dense( + small_binary, + "> + + result = EXGBoost.NIF.dmatrix_create_from_dense( + binary, + "> + + result = EXGBoost.NIF.dmatrix_create_from_dense( + binary, + "> + + # Little-endian should work + assert {:ok, _} = EXGBoost.NIF.dmatrix_create_from_dense( + binary, "> + + # Extremely large shape that would overflow + # Note: This might not fail if XGBoost itself fails first, + # but our validation should catch reasonable overflows + result = EXGBoost.NIF.dmatrix_create_from_dense( + binary, + "> + + # Integer should fail + assert {:error, _} = EXGBoost.NIF.dmatrix_create_from_dense( + binary, " -1.0}, Enum.into(opts, %{})) + |> Jason.encode!() + end end