├── .github ├── ISSUE_TEMPLATE │ ├── bug-report.md │ ├── documentation-request.md │ ├── feature-request.md │ ├── submit-question.md │ └── task.md ├── copy-pr-bot.yaml ├── release-drafter.yml └── workflows │ ├── check-base-branch.yaml │ ├── cpu-ci.yml │ ├── docs-preview-pr.yaml │ ├── docs-remove-stale-reviews.yaml │ ├── docs-sched-rebuild.yaml │ ├── gpu-ci.yml │ ├── lint.yaml │ ├── packages.yaml │ ├── postmerge-cpu.yml │ ├── postmerge-gpu.yml │ ├── release-drafter.yml │ ├── require-label.yaml │ ├── set-stable-branch.yaml │ └── triage.yml ├── .gitignore ├── .pre-commit-config.yaml ├── .prettierignore ├── .pylintrc ├── CLA.md ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── ci ├── ignore_codespell_words.txt ├── pr.gpu.Jenkinsfile ├── test_integration.sh └── test_unit.sh ├── conda └── recipes │ └── meta.yaml ├── docs ├── Makefile ├── README.md ├── make.bat └── source │ ├── _static │ ├── .gitkeep │ ├── NVIDIA-LogoBlack.svg │ ├── NVIDIA-LogoWhite.svg │ ├── css │ │ ├── custom.css │ │ └── versions.css │ ├── favicon.png │ └── js │ │ └── rtd-version-switcher.js │ ├── _templates │ ├── footer.html │ ├── layout.html │ ├── merlin-ecosystem.html │ └── versions.html │ ├── api.rst │ ├── conf.py │ ├── index.rst │ └── toc.yaml ├── merlin └── systems │ ├── __init__.py │ ├── _version.py │ ├── dag │ ├── __init__.py │ ├── ensemble.py │ ├── ops │ │ ├── __init__.py │ │ ├── compat.py │ │ ├── faiss.py │ │ ├── feast.py │ │ ├── fil.py │ │ ├── implicit.py │ │ ├── pytorch.py │ │ ├── session_filter.py │ │ ├── softmax_sampling.py │ │ ├── tensorflow.py │ │ ├── unroll_features.py │ │ └── workflow.py │ └── runtimes │ │ ├── __init__.py │ │ └── triton │ │ ├── __init__.py │ │ ├── ops │ │ ├── __init__.py │ │ ├── fil.py │ │ ├── operator.py │ │ ├── pytorch.py │ │ ├── tensorflow.py │ │ └── workflow.py │ │ └── runtime.py │ ├── model_registry.py │ ├── triton │ ├── __init__.py │ ├── conversions.py │ ├── export.py │ ├── models │ │ ├── __init__.py │ │ ├── executor_model.py │ │ ├── pytorch_model.py │ │ └── workflow_model.py │ └── utils.py │ └── workflow │ ├── __init__.py │ └── base.py ├── pyproject.toml ├── pytest.ini ├── requirements ├── base.txt ├── dev.txt ├── docs.txt ├── gpu.txt ├── test-cpu.txt ├── test-gpu.txt └── test.txt ├── setup.cfg ├── setup.py ├── tests ├── __init__.py ├── conftest.py ├── integration │ ├── feast │ │ └── test_int_feast.py │ ├── t4r │ │ └── test_pytorch_backend.py │ └── tf │ │ └── test_transformer_model.py ├── test_passing.py ├── unit │ ├── __init__.py │ └── systems │ │ ├── __init__.py │ │ ├── dag │ │ ├── ops │ │ │ ├── __init__.py │ │ │ └── test_ops.py │ │ ├── runtimes │ │ │ ├── __init__.py │ │ │ ├── local │ │ │ │ ├── __init__.py │ │ │ │ └── ops │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── fil │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── test_lightgbm.py │ │ │ │ │ ├── test_sklearn.py │ │ │ │ │ └── test_xgboost.py │ │ │ │ │ ├── nvtabular │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── test_ensemble.py │ │ │ │ │ ├── tensorflow │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── test_ensemble.py │ │ │ │ │ └── torch │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── test_op.py │ │ │ ├── local_runtime │ │ │ │ ├── __init__.py │ │ │ │ └── ops │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── torch │ │ │ │ │ └── __init__.py │ │ │ └── triton │ │ │ │ ├── __init__.py │ │ │ │ ├── ops │ │ │ │ ├── __init__.py │ │ │ │ ├── fil │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── test_lightgbm_triton.py │ │ │ │ │ ├── test_sklearn_triton.py │ │ │ │ │ └── test_xgboost_triton.py │ │ │ │ ├── torch │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── test_op.py │ │ │ │ └── workflow │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── test_ensemble.py │ │ │ │ │ └── test_op.py │ │ │ │ └── test_triton.py │ │ ├── test_ensemble.py │ │ ├── test_executors.py │ │ └── test_graph.py │ │ ├── ops │ │ ├── __init__.py │ │ ├── embedding_op.py │ │ ├── faiss │ │ │ └── test_executor.py │ │ ├── feast │ │ │ ├── __init__.py │ │ │ └── test_op.py │ │ ├── fil │ │ │ ├── __init__.py │ │ │ ├── test_ensemble.py │ │ │ ├── test_forest.py │ │ │ └── test_op.py │ │ ├── implicit │ │ │ ├── __init__.py │ │ │ ├── test_executor.py │ │ │ └── test_op.py │ │ ├── nvtabular │ │ │ └── __init__.py │ │ ├── padding_op.py │ │ ├── tf │ │ │ ├── __init__.py │ │ │ ├── test_ensemble.py │ │ │ └── test_op.py │ │ └── torch │ │ │ ├── __init__.py │ │ │ └── test_ensemble.py │ │ └── utils │ │ ├── __init__.py │ │ ├── ops.py │ │ ├── tf.py │ │ └── torch.py └── version │ └── test_version.py ├── tox.ini └── versioneer.py /.github/ISSUE_TEMPLATE/bug-report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "\U0001F41B Bug Report" 3 | about: Submit a bug report to help us improve Merlin Systems 4 | title: "[BUG]" 5 | labels: "status/needs-triage, bug" 6 | assignees: "" 7 | --- 8 | 9 | ### Bug description 10 | 11 | 12 | 13 | ### Steps/Code to reproduce bug 14 | 15 | 16 | 17 | 1. 18 | 2. 19 | 3. 20 | 21 | ### Expected behavior 22 | 23 | 24 | 25 | ### Environment details 26 | 27 | - Merlin version: 28 | - Platform: 29 | - Python version: 30 | - PyTorch version (GPU?): 31 | - Tensorflow version (GPU?): 32 | 33 | ### Additional context 34 | 35 | 36 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/documentation-request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Documentation request 3 | about: Report incorrect or needed documentation 4 | title: "[DOC]" 5 | labels: "status/needs-triage, area/documentation" 6 | assignees: "" 7 | --- 8 | 9 | ## Report incorrect documentation 10 | 11 | ### Location of incorrect documentation 12 | 13 | 14 | 15 | ### Describe the problems or issues found in the documentation 16 | 17 | 18 | 19 | ### Steps taken to verify documentation is incorrect 20 | 21 | 22 | 23 | ### Suggested fix for documentation 24 | 25 | 26 | 27 | --- 28 | 29 | ## Report needed documentation 30 | 31 | ### Report needed documentation 32 | 33 | 34 | 35 | ### Describe the documentation you'd like 36 | 37 | 38 | 39 | ### Steps taken to search for needed documentation\*\* 40 | 41 | 42 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature-request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "\U0001F680 Feature request" 3 | about: Submit a proposal/request for a new Merlin Systems feature 4 | title: "[FEA]" 5 | labels: "status/needs-triage, kind/feature-request" 6 | assignees: "" 7 | --- 8 | 9 | # 🚀 Feature request 10 | 11 | 13 | 14 | ## Motivation 15 | 16 | 19 | 20 | ## Your contribution 21 | 22 | 25 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/submit-question.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "❓ Questions & Help" 3 | about: Ask a general question about Merlin Systems 4 | title: "[QST]" 5 | labels: "status/needs-triage, kind/question" 6 | assignees: "" 7 | --- 8 | 9 | # ❓ Questions & Help 10 | 11 | ## Details 12 | 13 | 14 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/task.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Task 3 | about: A general task that we're tracking in Github 4 | title: "[Task]" 5 | labels: "" 6 | assignees: "" 7 | --- 8 | 9 | ### Description 10 | 11 | **Additional context** 12 | Add any other context, code examples, or references to existing implementations about the task here. 13 | -------------------------------------------------------------------------------- /.github/copy-pr-bot.yaml: -------------------------------------------------------------------------------- 1 | # Configuration file for `copy-pr-bot` GitHub App 2 | # https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/ 3 | 4 | enabled: true 5 | -------------------------------------------------------------------------------- /.github/release-drafter.yml: -------------------------------------------------------------------------------- 1 | categories: 2 | - title: "⚠ Breaking Changes" 3 | labels: 4 | - "breaking" 5 | - title: "🐜 Bug Fixes" 6 | labels: 7 | - "bug" 8 | - title: "🚀 Features" 9 | labels: 10 | - "feature" 11 | - "enhancement" 12 | - title: "📄 Documentation" 13 | labels: 14 | - "documentation" 15 | - "examples" 16 | - title: "🔧 Maintenance" 17 | labels: 18 | - "build" 19 | - "dependencies" 20 | - "chore" 21 | - "ci" 22 | change-template: "- $TITLE @$AUTHOR (#$NUMBER)" 23 | exclude-labels: 24 | - "skip-changelog" 25 | template: | 26 | ## What’s Changed 27 | 28 | $CHANGES 29 | -------------------------------------------------------------------------------- /.github/workflows/check-base-branch.yaml: -------------------------------------------------------------------------------- 1 | name: Require Development Base Branch 2 | 3 | on: 4 | pull_request: 5 | types: [synchronize, opened, reopened, labeled, unlabeled] 6 | 7 | jobs: 8 | check: 9 | uses: NVIDIA-Merlin/.github/.github/workflows/check-base-branch.yml@main 10 | -------------------------------------------------------------------------------- /.github/workflows/cpu-ci.yml: -------------------------------------------------------------------------------- 1 | name: CPU CI 2 | 3 | on: 4 | workflow_dispatch: 5 | pull_request: 6 | branches: [main] 7 | types: [opened, synchronize, reopened] 8 | 9 | concurrency: 10 | group: ${{ github.workflow }}-${{ github.ref }} 11 | cancel-in-progress: true 12 | 13 | jobs: 14 | cpu-ci-premerge: 15 | runs-on: ${{ matrix.os }} 16 | strategy: 17 | matrix: 18 | python-version: [3.8] 19 | os: [ubuntu-latest] 20 | 21 | steps: 22 | - uses: actions/checkout@v3 23 | with: 24 | fetch-depth: 0 25 | - name: Set up Python ${{ matrix.python-version }} 26 | uses: actions/setup-python@v4 27 | with: 28 | python-version: ${{ matrix.python-version }} 29 | - name: Install Ubuntu packages 30 | run: | 31 | sudo apt-get update -y 32 | sudo apt-get install -y protobuf-compiler 33 | - name: Install and upgrade python packages 34 | run: | 35 | python -m pip install --upgrade pip setuptools==59.4.0 wheel tox 36 | - name: Get Branch name 37 | id: get-branch-name 38 | uses: NVIDIA-Merlin/.github/actions/branch-name@main 39 | - name: Run tests 40 | run: | 41 | branch="${{ steps.get-branch-name.outputs.branch }}" 42 | tox -e test-cpu -- $branch 43 | - name: Building docs 44 | run: | 45 | tox -e docs 46 | - name: Upload HTML 47 | uses: actions/upload-artifact@v3 48 | with: 49 | name: html-build-artifact 50 | path: docs/build/html 51 | if-no-files-found: error 52 | retention-days: 1 53 | - name: Store PR information 54 | run: | 55 | mkdir ./pr 56 | echo ${{ github.event.number }} > ./pr/pr.txt 57 | echo ${{ github.event.pull_request.merged }} > ./pr/merged.txt 58 | echo ${{ github.event.action }} > ./pr/action.txt 59 | - name: Upload PR information 60 | uses: actions/upload-artifact@v3 61 | with: 62 | name: pr 63 | path: pr/ 64 | -------------------------------------------------------------------------------- /.github/workflows/docs-preview-pr.yaml: -------------------------------------------------------------------------------- 1 | name: docs-preview-pr 2 | 3 | on: 4 | workflow_run: 5 | workflows: [CPU CI] 6 | types: [completed] 7 | 8 | concurrency: 9 | group: ${{ github.workflow }}-${{ github.ref }} 10 | cancel-in-progress: true 11 | 12 | env: 13 | WF_ID: ${{ github.event.workflow_run.id }} 14 | 15 | jobs: 16 | preview: 17 | uses: nvidia-merlin/.github/.github/workflows/docs-preview-pr-common.yaml@main 18 | -------------------------------------------------------------------------------- /.github/workflows/docs-remove-stale-reviews.yaml: -------------------------------------------------------------------------------- 1 | name: docs-remove-stale-reviews 2 | 3 | on: 4 | schedule: 5 | # 42 minutes after 0:00 UTC on Sundays 6 | - cron: "42 0 * * 0" 7 | workflow_dispatch: 8 | 9 | jobs: 10 | remove: 11 | uses: nvidia-merlin/.github/.github/workflows/docs-remove-stale-reviews-common.yaml@main 12 | -------------------------------------------------------------------------------- /.github/workflows/docs-sched-rebuild.yaml: -------------------------------------------------------------------------------- 1 | name: docs-sched-rebuild 2 | 3 | on: 4 | push: 5 | branches: [main] 6 | tags: 7 | - "v[0-9]+.[0-9]+.[0-9]+" 8 | workflow_dispatch: 9 | 10 | jobs: 11 | build: 12 | runs-on: [ubuntu-latest] 13 | 14 | steps: 15 | - uses: actions/checkout@v3 16 | with: 17 | fetch-depth: 0 18 | - name: Set up Python 3.9 19 | uses: actions/setup-python@v4 20 | with: 21 | python-version: 3.9 22 | - name: Install Ubuntu packages 23 | run: | 24 | sudo apt-get update -y 25 | sudo apt-get install -y protobuf-compiler 26 | - name: Install dependencies 27 | run: | 28 | python -m pip install --upgrade pip setuptools==59.4.0 wheel tox 29 | - name: Setup local branches for docs build 30 | run: | 31 | git branch --track main origin/main || true 32 | git branch --track stable origin/stable || true 33 | - name: Building docs (multiversion) 34 | run: | 35 | tox -e docs-multi 36 | - name: Delete unnecessary files 37 | run: | 38 | find docs/build -name .doctrees -prune -exec rm -rf {} \; 39 | find docs/build -name .buildinfo -exec rm {} \; 40 | - name: Upload HTML 41 | uses: actions/upload-artifact@v3 42 | with: 43 | name: html-build-artifact 44 | path: docs/build/html 45 | if-no-files-found: error 46 | retention-days: 1 47 | 48 | # Identify the dir for the HTML. 49 | store-html: 50 | needs: [build] 51 | runs-on: ubuntu-latest 52 | steps: 53 | - uses: actions/checkout@v3 54 | with: 55 | ref: "gh-pages" 56 | - name: Initialize Git configuration 57 | run: | 58 | git config user.name docs-sched-rebuild 59 | git config user.email do-not-send-@github.com 60 | - name: Download artifacts 61 | uses: actions/download-artifact@v3 62 | with: 63 | name: html-build-artifact 64 | - name: Copy HTML directories 65 | run: | 66 | ls -asl 67 | for i in `ls -d *` 68 | do 69 | echo "Git adding ${i}" 70 | git add "${i}" 71 | done 72 | - name: Check or create dot-no-jekyll file 73 | run: | 74 | if [ -f ".nojekyll" ]; then 75 | echo "The dot-no-jekyll file already exists." 76 | exit 0 77 | fi 78 | touch .nojekyll 79 | git add .nojekyll 80 | - name: Check or create redirect page 81 | env: 82 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 83 | run: | 84 | resp=$(grep 'http-equiv="refresh"' index.html 2>/dev/null) || true 85 | if [ -n "${resp}" ]; then 86 | echo "The redirect file already exists." 87 | exit 0 88 | fi 89 | def_branch="stable" 90 | html_url=$(gh api "repos/${GITHUB_REPOSITORY}/pages" --jq ".html_url") 91 | cat > index.html << EOF 92 | 93 | 94 | 95 | Redirect to documentation 96 | 97 | 99 | 104 | 105 | 106 |

Please follow the link to the 107 | ${def_branch}' branch documentation.

108 | 109 | 110 | EOF 111 | git add index.html 112 | - name: Commit changes to the GitHub Pages branch 113 | run: | 114 | git status 115 | if git commit -m 'Pushing changes to GitHub Pages.'; then 116 | git push -f 117 | else 118 | echo "Nothing changed." 119 | fi 120 | -------------------------------------------------------------------------------- /.github/workflows/gpu-ci.yml: -------------------------------------------------------------------------------- 1 | name: GPU CI 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: 7 | - main 8 | - "pull-request/[0-9]+" 9 | tags: 10 | - "v[0-9]+.[0-9]+.[0-9]+" 11 | 12 | jobs: 13 | gpu-ci-premerge: 14 | runs-on: linux-amd64-gpu-p100-latest-1 15 | container: 16 | image: nvcr.io/nvstaging/merlin/merlin-ci-runner:latest 17 | env: 18 | NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }} 19 | options: --shm-size=1G 20 | credentials: 21 | username: $oauthtoken 22 | password: ${{ secrets.NGC_TOKEN }} 23 | 24 | steps: 25 | - uses: actions/checkout@v3 26 | with: 27 | fetch-depth: 0 28 | - name: Run tests 29 | run: | 30 | ref_type=${{ github.ref_type }} 31 | branch=main 32 | if [[ $ref_type == "tag"* ]] 33 | then 34 | raw=$(git branch -r --contains ${{ github.ref_name }}) 35 | branch=${raw/origin\/} 36 | fi 37 | tox -e test-gpu -- $branch 38 | -------------------------------------------------------------------------------- /.github/workflows/lint.yaml: -------------------------------------------------------------------------------- 1 | name: lint 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: [main] 7 | tags: 8 | - "v[0-9]+.[0-9]+.[0-9]+" 9 | pull_request: 10 | branches: [main] 11 | types: [opened, synchronize, reopened] 12 | 13 | concurrency: 14 | group: ${{ github.workflow }}-${{ github.ref }} 15 | cancel-in-progress: true 16 | 17 | jobs: 18 | lint: 19 | runs-on: ubuntu-latest 20 | strategy: 21 | matrix: 22 | python-version: [3.8] 23 | steps: 24 | - uses: actions/checkout@v3 25 | - uses: actions/setup-python@v4 26 | with: 27 | python-version: "3.8" 28 | - uses: pre-commit/action@v3.0.0 29 | -------------------------------------------------------------------------------- /.github/workflows/packages.yaml: -------------------------------------------------------------------------------- 1 | name: Packages 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: 7 | - main 8 | - release* 9 | tags: 10 | - v* 11 | pull_request: 12 | branches: [main] 13 | 14 | concurrency: 15 | group: ${{ github.workflow }}-${{ github.ref }} 16 | cancel-in-progress: true 17 | 18 | jobs: 19 | build-pypi: 20 | name: Build PyPI Package 21 | runs-on: ubuntu-latest 22 | steps: 23 | - uses: actions/checkout@v3 24 | - name: Set up Python 3.9 25 | uses: actions/setup-python@v4 26 | with: 27 | python-version: 3.9 28 | - name: Install and upgrade python packages 29 | run: | 30 | python -m pip install --upgrade pip setuptools==59.4.0 wheel 31 | - name: Generate package for pypi 32 | run: | 33 | python setup.py sdist 34 | - name: Upload pypi artifacts to github 35 | uses: actions/upload-artifact@v3 36 | with: 37 | name: dist 38 | path: dist 39 | 40 | build-conda: 41 | name: Build Conda Package 42 | runs-on: ubuntu-latest 43 | steps: 44 | - uses: actions/checkout@v3 45 | with: 46 | fetch-depth: 0 47 | - name: Set up Python 3.9 48 | uses: actions/setup-python@v4 49 | with: 50 | python-version: 3.9 51 | - uses: conda-incubator/setup-miniconda@v2 52 | with: 53 | auto-update-conda: true 54 | - name: Generate package for conda 55 | id: conda_build 56 | run: | 57 | echo "conda pkgs dir $CONDA_PKGS_DIRS" 58 | conda install -c conda-forge mamba 59 | mamba install -c conda-forge conda-build boa conda-verify 60 | conda mambabuild . -c defaults -c conda-forge -c numba -c rapidsai -c nvidia --output-folder ./conda_packages 61 | conda_package=$(find ./conda_packages/ -name "*.tar.bz2") 62 | export CONDA_PACKAGE=$conda_package 63 | echo "conda_package : $conda_package" 64 | echo "conda_package=$conda_package" >> "$GITHUB_OUTPUT" 65 | - name: Upload conda artifacts to github 66 | uses: actions/upload-artifact@v3 67 | with: 68 | name: conda 69 | path: ${{ steps.conda_build.outputs.conda_package }} 70 | 71 | release-pypi: 72 | name: Release PyPI Package 73 | runs-on: ubuntu-latest 74 | if: "startsWith(github.ref, 'refs/tags/')" 75 | needs: [build-pypi] 76 | steps: 77 | - uses: actions/download-artifact@v3 78 | with: 79 | name: dist 80 | path: dist 81 | - name: Create GitHub Release 82 | uses: fnkr/github-action-ghr@v1.3 83 | env: 84 | GHR_PATH: ./dist 85 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 86 | - uses: actions/setup-python@v4 87 | with: 88 | python-version: 3.9 89 | - name: Push to PyPI 90 | env: 91 | TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} 92 | TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} 93 | run: | 94 | pip install --upgrade wheel pip setuptools twine 95 | twine upload dist/* 96 | 97 | release-conda: 98 | name: Release Conda Package 99 | runs-on: ubuntu-latest 100 | if: "startsWith(github.ref, 'refs/tags/')" 101 | needs: [build-conda] 102 | steps: 103 | - uses: actions/setup-python@v2 104 | with: 105 | python-version: 3.9 106 | - uses: actions/download-artifact@v3 107 | with: 108 | name: conda 109 | path: conda 110 | - uses: conda-incubator/setup-miniconda@v2 111 | with: 112 | auto-update-conda: true 113 | - name: Install conda dependencies 114 | shell: bash -l {0} 115 | run: | 116 | conda install -y anaconda-client conda-build 117 | - name: Push to anaconda 118 | shell: bash -l {0} 119 | env: 120 | ANACONDA_TOKEN: ${{ secrets.ANACONDA_TOKEN }} 121 | run: | 122 | anaconda -t $ANACONDA_TOKEN upload -u nvidia conda/*.tar.bz2 123 | -------------------------------------------------------------------------------- /.github/workflows/postmerge-cpu.yml: -------------------------------------------------------------------------------- 1 | name: CPU CI (Post-merge) 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: [main] 7 | tags: 8 | - "v[0-9]+.[0-9]+.[0-9]+" 9 | 10 | jobs: 11 | cpu-ci-postmerge: 12 | runs-on: ${{ matrix.os }} 13 | strategy: 14 | matrix: 15 | python-version: [3.8] 16 | os: [ubuntu-latest] 17 | 18 | steps: 19 | - uses: actions/checkout@v3 20 | with: 21 | fetch-depth: 0 22 | - name: Set up Python ${{ matrix.python-version }} 23 | uses: actions/setup-python@v4 24 | with: 25 | python-version: ${{ matrix.python-version }} 26 | - name: Install Ubuntu packages 27 | run: | 28 | sudo apt-get update -y 29 | sudo apt-get install -y protobuf-compiler 30 | - name: Install and upgrade python packages 31 | run: | 32 | python -m pip install --upgrade pip setuptools==59.4.0 wheel tox 33 | - name: Get Branch name 34 | id: get-branch-name 35 | uses: NVIDIA-Merlin/.github/actions/branch-name@main 36 | - name: Run tests 37 | run: | 38 | branch="${{ steps.get-branch-name.outputs.branch }}" 39 | tox -e test-cpu-postmerge -- $branch 40 | - name: Building docs 41 | run: | 42 | tox -e docs 43 | - name: Upload HTML 44 | uses: actions/upload-artifact@v3 45 | with: 46 | name: html-build-artifact 47 | path: docs/build/html 48 | if-no-files-found: error 49 | retention-days: 1 50 | - name: Store PR information 51 | run: | 52 | mkdir ./pr 53 | echo ${{ github.event.number }} > ./pr/pr.txt 54 | echo ${{ github.event.pull_request.merged }} > ./pr/merged.txt 55 | echo ${{ github.event.action }} > ./pr/action.txt 56 | - name: Upload PR information 57 | uses: actions/upload-artifact@v3 58 | with: 59 | name: pr 60 | path: pr/ 61 | 62 | # If failures occur, notify in slack. 63 | - name: Notify Failure in Slack 64 | id: slack 65 | uses: slackapi/slack-github-action@v1.23.0 66 | if: ${{ failure() && github.ref == 'refs/heads/main' }} 67 | with: 68 | # This data can be any valid JSON from a previous step in the GitHub Action 69 | payload: | 70 | { 71 | "message": "GitHub Action job failure on github.ref `${{ github.ref }}`", 72 | "pull_request": "${{ github.event.pull_request.html_url || github.event.head_commit.url }}", 73 | "workflow": "${{ github.workflow }}", 74 | "logs_url": "${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}" 75 | } 76 | env: 77 | PR_NUMBER: ${{ github.event.number }} 78 | SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }} 79 | -------------------------------------------------------------------------------- /.github/workflows/postmerge-gpu.yml: -------------------------------------------------------------------------------- 1 | name: GPU CI (Post-merge) 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: [main] 7 | tags: 8 | - "v[0-9]+.[0-9]+.[0-9]+" 9 | 10 | jobs: 11 | gpu-ci-postmerge: 12 | runs-on: 1GPU 13 | 14 | steps: 15 | - uses: actions/checkout@v3 16 | with: 17 | fetch-depth: 0 18 | - name: Run tests 19 | run: | 20 | ref_type=${{ github.ref_type }} 21 | branch=main 22 | if [[ $ref_type == "tag"* ]] 23 | then 24 | raw=$(git branch -r --contains ${{ github.ref_name }}) 25 | branch=${raw/origin\/} 26 | fi 27 | cd ${{ github.workspace }}; tox -e test-gpu-postmerge -- $branch 28 | 29 | # If failures occur, notify in slack. 30 | - name: Notify Failure in Slack 31 | id: slack 32 | uses: slackapi/slack-github-action@v1.23.0 33 | if: ${{ failure() && github.ref == 'refs/heads/main' }} 34 | with: 35 | # This data can be any valid JSON from a previous step in the GitHub Action 36 | payload: | 37 | { 38 | "message": "GitHub Action job failure on github.ref `${{ github.ref }}`", 39 | "pull_request": "${{ github.event.pull_request.html_url || github.event.head_commit.url }}", 40 | "workflow": "${{ github.workflow }}", 41 | "logs_url": "${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}" 42 | } 43 | env: 44 | PR_NUMBER: ${{ github.event.number }} 45 | SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }} 46 | -------------------------------------------------------------------------------- /.github/workflows/release-drafter.yml: -------------------------------------------------------------------------------- 1 | name: release-drafter 2 | 3 | on: 4 | push: 5 | # trigger on tags only 6 | tags: 7 | - "v[0-9]+.[0-9]+.[0-9]+" 8 | 9 | workflow_dispatch: 10 | 11 | jobs: 12 | update_release_draft: 13 | uses: nvidia-merlin/.github/.github/workflows/release-drafter-common.yaml@main 14 | -------------------------------------------------------------------------------- /.github/workflows/require-label.yaml: -------------------------------------------------------------------------------- 1 | name: Require PR Labels 2 | 3 | on: 4 | pull_request: 5 | types: [synchronize, opened, reopened, labeled, unlabeled] 6 | 7 | jobs: 8 | check-labels: 9 | uses: nvidia-merlin/.github/.github/workflows/require-label.yaml@main 10 | -------------------------------------------------------------------------------- /.github/workflows/set-stable-branch.yaml: -------------------------------------------------------------------------------- 1 | name: Set Stable Branch 2 | 3 | on: 4 | workflow_dispatch: 5 | release: 6 | types: [published, deleted] 7 | 8 | jobs: 9 | set-stable-branch: 10 | uses: NVIDIA-Merlin/.github/.github/workflows/set-stable-branch.yaml@main 11 | -------------------------------------------------------------------------------- /.github/workflows/triage.yml: -------------------------------------------------------------------------------- 1 | name: triage_issues 2 | on: 3 | issues: 4 | types: [opened, reopened] 5 | 6 | jobs: 7 | triage_issue: 8 | uses: nvidia-merlin/.github/.github/workflows/triage.yaml@main 9 | secrets: 10 | TRIAGE_APP_ID: ${{ secrets.TRIAGE_APP_ID }} 11 | TRIAGE_APP_PEM: ${{ secrets.TRIAGE_APP_PEM }} 12 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | **/.ipynb_checkpoints/* 2 | /.*_checkpoints/ 3 | .ipynb_checkpoints/* 4 | */.jupyter/* 5 | */.local/* 6 | 7 | cache/* 8 | data*/* 9 | *.parquet 10 | *.orc 11 | *.csv 12 | 13 | 14 | # Byte-compiled / optimized / DLL files 15 | __pycache__/ 16 | *.py[cod] 17 | *$py.class 18 | 19 | # C extensions 20 | *.so 21 | 22 | # Distribution / packaging 23 | .Python 24 | build/ 25 | develop-eggs/ 26 | dist/ 27 | downloads/ 28 | eggs/ 29 | .eggs/ 30 | lib/ 31 | lib64/ 32 | parts/ 33 | sdist/ 34 | var/ 35 | wheels/ 36 | share/python-wheels/ 37 | *.egg-info/ 38 | .installed.cfg 39 | *.egg 40 | MANIFEST 41 | 42 | # Installer logs 43 | pip-log.txt 44 | pip-delete-this-directory.txt 45 | 46 | # Unit test / coverage reports 47 | htmlcov/ 48 | .tox/ 49 | .nox/ 50 | .coverage 51 | .coverage.* 52 | .cache 53 | nosetests.xml 54 | coverage.xml 55 | *.cover 56 | *.py,cover 57 | .hypothesis/ 58 | .pytest_cache/ 59 | cover/ 60 | 61 | # Sphinx documentation 62 | docs/build/ 63 | docs/source/LICENSE 64 | docs/source/generated/ 65 | docs/source/README.md 66 | docs/source/examples/ 67 | 68 | # IPython 69 | profile_default/ 70 | ipython_config.py 71 | 72 | # mypy 73 | .mypy_cache/ 74 | .dmypy.json 75 | dmypy.json 76 | 77 | # PyCharm 78 | .idea 79 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/timothycrosley/isort 3 | rev: 5.12.0 4 | hooks: 5 | - id: isort 6 | additional_dependencies: [toml] 7 | exclude: examples/.* 8 | - repo: https://github.com/python/black 9 | rev: 22.12.0 10 | hooks: 11 | - id: black 12 | - repo: https://github.com/pycqa/flake8 13 | rev: 6.0.0 14 | hooks: 15 | - id: flake8 16 | - repo: https://github.com/pycqa/pylint 17 | rev: v2.15.8 18 | hooks: 19 | - id: pylint 20 | - repo: https://github.com/pre-commit/mirrors-prettier 21 | rev: v2.7.1 22 | hooks: 23 | - id: prettier 24 | types_or: [yaml, markdown] 25 | - repo: https://github.com/econchick/interrogate 26 | rev: 1.5.0 27 | hooks: 28 | - id: interrogate 29 | exclude: ^(docs|tests|setup.py|versioneer.py) 30 | args: [--config=pyproject.toml] 31 | - repo: https://github.com/codespell-project/codespell 32 | rev: v2.2.2 33 | hooks: 34 | - id: codespell 35 | - repo: https://github.com/PyCQA/bandit 36 | rev: 1.7.4 37 | hooks: 38 | - id: bandit 39 | args: [--verbose, -ll, -x, tests, examples, bench] 40 | -------------------------------------------------------------------------------- /.prettierignore: -------------------------------------------------------------------------------- 1 | build/ 2 | conda/ 3 | -------------------------------------------------------------------------------- /.pylintrc: -------------------------------------------------------------------------------- 1 | [MASTER] 2 | 3 | ignore-patterns=_version.py,model_config_pb2.py 4 | 5 | extension-pkg-allow-list=hugectr,nvtabular_cpp 6 | 7 | [MESSAGES CONTROL] 8 | disable=fixme, 9 | # docstrings aren't required (yet). 10 | missing-function-docstring, 11 | missing-module-docstring, 12 | missing-class-docstring, 13 | 14 | # formatting checks that we're handling with black/isort 15 | wrong-import-order, 16 | wrong-import-position, 17 | ungrouped-imports, 18 | line-too-long, 19 | superfluous-parens, 20 | trailing-whitespace, 21 | 22 | # we'll probably never enable these checks 23 | invalid-name, 24 | import-error, 25 | 26 | # disable code-complexity checks for now 27 | # TODO: should we configure the thresholds for these rather than just disable? 28 | too-many-function-args, 29 | too-many-instance-attributes, 30 | too-many-locals, 31 | too-many-branches, 32 | too-many-nested-blocks, 33 | too-many-statements, 34 | too-many-arguments, 35 | too-many-return-statements, 36 | too-many-lines, 37 | too-few-public-methods, 38 | 39 | # many of these checks would be great to include at some point, but would 40 | # require some changes to our codebase 41 | useless-return, 42 | protected-access, 43 | arguments-differ, 44 | unused-argument, 45 | unused-variable, 46 | abstract-method, 47 | no-name-in-module, 48 | attribute-defined-outside-init, 49 | redefined-outer-name, 50 | import-outside-toplevel, 51 | no-else-continue, 52 | no-else-return, 53 | no-else-raise, 54 | no-member, 55 | super-with-arguments, 56 | unsupported-assignment-operation, 57 | inconsistent-return-statements, 58 | duplicate-string-formatting-argument, 59 | len-as-condition, 60 | cyclic-import, 61 | 62 | # producing false positives 63 | unexpected-keyword-arg, 64 | not-an-iterable, 65 | unsubscriptable-object, 66 | signature-differs 67 | 68 | [SIMILARITIES] 69 | min-similarity-lines=50 70 | ignore-comments=yes 71 | ignore-docstrings=yes 72 | ignore-imports=yes 73 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to Merlin Systems 2 | 3 | If you are interested in contributing to Merlin Systems, your contributions will fall 4 | into three categories: 5 | 6 | 1. You want to report a bug, feature request, or documentation issue: 7 | - File an [issue](https://github.com/NVIDIA-Merlin/systems/issues/new/choose) 8 | and describe what you encountered or what you want to see changed. 9 | - The NVIDIA-Merlin team evaluates the issues and triages them, scheduling 10 | them for a release. If you believe the issue needs priority attention, 11 | comment on the issue to notify the team. 12 | 2. You want to propose a new feature and implement it: 13 | - Post about your intended feature to discuss the design and 14 | implementation with the NVIDIA-Merlin team. 15 | - Once we agree that the plan looks good, go ahead and implement it, using 16 | the [code contributions](#code-contributions) guide below. 17 | 3. You want to implement a feature or bug-fix for an outstanding issue: 18 | - Follow the [code contributions](#code-contributions) guide below. 19 | - If you need more context on a particular issue, please ask and the 20 | NVIDIA-Merlin team will provide the context. 21 | 22 | ## Code contributions 23 | 24 | ### Your first issue 25 | 26 | 1. Read the project's [README.md](https://github.com/NVIDIA-Merlin/systems/blob/stable/README.md) 27 | to learn how to setup the development environment. 28 | 2. Find an issue to work on. The best way is to look for the [good first issue](https://github.com/NVIDIA-Merlin/systems/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) 29 | or [help wanted](https://github.com/NVIDIA-Merlin/systems/issues?q=is%3Aissue+is%3Aopen+label%3A%22help+wanted%22) labels. 30 | 3. Comment on the issue to say that you are going to work on it. 31 | 4. Code! Make sure to update unit tests! 32 | 5. When done, [create your pull request](https://github.com/NVIDIA-Merlin/systems/compare). 33 | 6. Verify that CI passes all [status checks](https://help.github.com/articles/about-status-checks/). Fix if needed. 34 | 7. Wait for other developers to review your code and update code as needed. 35 | 8. After your pull request is reviewed and approved, an NVIDIA-Merlin team member merges it. 36 | 37 | Remember, if you are unsure about anything, don't hesitate to comment on issues 38 | and ask for clarifications! 39 | 40 | ### Seasoned developers 41 | 42 | After you have your feet wet and are comfortable with the code, you 43 | can look at the prioritized issues of our next release in our [project boards](https://github.com/NVIDIA-Merlin/systems/projects). 44 | 45 | > **Pro Tip:** Always look at the release board with the highest number for 46 | > issues to work on. This is where the NVIDIA-Merlin developers also focus their efforts. 47 | 48 | Look at the unassigned issues, and find an issue you are comfortable with 49 | contributing to. Start with _Step 3_ from above, commenting on the issue to let 50 | others know that you are working on it. If you have any questions related to the 51 | implementation of the issue, ask them in the issue instead of the PR. 52 | 53 | ## Label your PRs 54 | 55 | This repository uses the release-drafter action to draft and create our change log. 56 | 57 | Please add one of the following labels to your PR to specify the type of contribution 58 | and help categorize the PR in our change log: 59 | 60 | - `breaking` -- The PR creates a breaking change to the API. 61 | - `bug` -- The PR fixes a problem with the code. 62 | - `feature` or `enhancement` -- The PR introduces a backward-compatible feature. 63 | - `documentation` or `examples` -- The PR is an addition or update to documentation. 64 | - `build`, `dependencies`, `chore`, or `ci` -- The PR is related to maintaining the 65 | repository or the project. 66 | 67 | By default, an unlabeled PR is listed at the top of the change log and is not 68 | grouped under a heading like _Features_ that groups similar PRs. 69 | Labeling the PRs so we can categorize them is preferred. 70 | 71 | If, for some reason, you do not believe your PR should be included in the change 72 | log, you can add the `skip-changelog` label. 73 | This label excludes the PR from the change log. 74 | 75 | For more information, see `.github/release-drafter.yml` in the repository 76 | or go to . 77 | 78 | ## Attribution 79 | 80 | Portions adopted from https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md 81 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include versioneer.py 2 | recursive-include merlin *.py *.proto 3 | include requirements/*.txt 4 | 5 | include merlin/core/_version.py 6 | -------------------------------------------------------------------------------- /ci/ignore_codespell_words.txt: -------------------------------------------------------------------------------- 1 | te 2 | coo 3 | ser 4 | fo 5 | ot 6 | lik 7 | usera 8 | als 9 | -------------------------------------------------------------------------------- /ci/pr.gpu.Jenkinsfile: -------------------------------------------------------------------------------- 1 | pipeline { 2 | agent { 3 | docker { 4 | image 'nvcr.io/nvstaging/merlin/merlin-ci-runner-wrapper' 5 | label 'merlin_gpu_gcp || merlin_gpu' 6 | registryCredentialsId 'jawe-nvcr-io' 7 | registryUrl 'https://nvcr.io' 8 | args "--runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all --shm-size '256m'" 9 | } 10 | } 11 | 12 | options { 13 | buildDiscarder(logRotator(numToKeepStr: '10')) 14 | ansiColor('xterm') 15 | disableConcurrentBuilds(abortPrevious: true) 16 | } 17 | 18 | stages { 19 | stage("test-gpu") { 20 | options { 21 | timeout(time: 60, unit: 'MINUTES', activity: false) 22 | } 23 | steps { 24 | sh """#!/bin/bash 25 | set -e 26 | printenv 27 | 28 | rm -rf $HOME/.cudf/ 29 | export TF_MEMORY_ALLOCATION="0.1" 30 | export CUDA_VISIBLE_DEVICES=2,3 31 | export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION='python' 32 | export MKL_SERVICE_FORCE_INTEL=1 33 | export NR_USER=true 34 | 35 | echo PYTHONPATH=\$PYTHONPATH:\$(pwd) 36 | 37 | tox -re test-gpu 38 | """ 39 | } 40 | } 41 | } 42 | } -------------------------------------------------------------------------------- /ci/test_integration.sh: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2023, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | #!/bin/bash 18 | set -e 19 | 20 | # Call this script with: 21 | # 1. Name of container as first parameter 22 | # [merlin-hugectr, merlin-tensorflow, merlin-pytorch] 23 | # 24 | # 2. Devices to use: 25 | # [0; 0,1; 0,1,..,n-1] 26 | 27 | cd /systems/ 28 | 29 | container=$1 30 | devices=$2 31 | 32 | pip install feast==0.31 33 | CUDA_VISIBLE_DEVICES="$devices" TF_GPU_ALLOCATOR=cuda_malloc_async python -m pytest -rxs tests/integration 34 | -------------------------------------------------------------------------------- /ci/test_unit.sh: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | #!/bin/bash 18 | set -e 19 | 20 | # Run tests 21 | TF_GPU_ALLOCATOR=cuda_malloc_async pytest -rsx tests/unit tests/notebook 22 | 23 | 24 | -------------------------------------------------------------------------------- /conda/recipes/meta.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2023, NVIDIA CORPORATION. 2 | 3 | # Usage: 4 | # conda build . -c defaults -c conda-forge -c numba -c rapidsai 5 | 6 | {% set version = environ.get('GIT_DESCRIBE_TAG', '0.1').lstrip('v') + environ.get('VERSION_SUFFIX', '') %} 7 | {% set git_revision_count=environ.get('GIT_DESCRIBE_NUMBER', 0) %} 8 | {% set setup_py_data = load_setup_py_data() %} 9 | 10 | package: 11 | name: merlin-systems 12 | version: {{ version }} 13 | 14 | source: 15 | path: ../../ 16 | 17 | build: 18 | number: {{ git_revision_count }} 19 | noarch: python 20 | script: python -m pip install . -vvv 21 | 22 | requirements: 23 | build: 24 | - python 25 | - setuptools 26 | run: 27 | - python 28 | {% for req in setup_py_data.get('install_requires', []) %} 29 | # the treelite_runtime pip package is included in the treelite conda package 30 | # and not available to install separately 31 | {% if not req.startswith("treelite_runtime") %} 32 | - {{ req }} 33 | {% endif %} 34 | {% endfor %} 35 | 36 | about: 37 | home: https://github.com/NVIDIA-Merlin/systems 38 | license_file: LICENSE 39 | summary: Merlin Systems provides tools for combining recommendation models with other elements of production recommender systems (like feature stores, nearest neighbor search, and exploration strategies) into end-to-end recommendation pipelines that can be served with Triton Inference Server. 40 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | clean: 16 | rm -rf build source/generated source/README.md source/examples 17 | 18 | 19 | .PHONY: help Makefile 20 | 21 | # Catch-all target: route all unknown targets to Sphinx using the new 22 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 23 | %: Makefile 24 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 25 | -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- 1 | # Documentation 2 | 3 | This folder contains the scripts necessary to build documentation. You can 4 | find the documentation at . 5 | 6 | ## Contributing to Docs 7 | 8 | You build the documentation with the `tox` command and specify the `docs` environment. 9 | The following steps are one way of many to build the documentation before opening a merge request. 10 | 11 | 1. Create a virtual environment: 12 | 13 | ```shell 14 | python -m venv .venv 15 | ``` 16 | 17 | 1. Activate the virtual environment: 18 | 19 | ```shell 20 | source .venv/bin/activate 21 | ``` 22 | 23 | 1. Install tox in the virtual environment: 24 | 25 | ```shell 26 | python -m pip install --upgrade pip 27 | python -m pip install tox 28 | ``` 29 | 30 | 1. Build the documentation with tox: 31 | 32 | ```shell 33 | tox -e docs 34 | ``` 35 | 36 | These steps run Sphinx in your shell and create HTML in the `docs/build/html/` 37 | directory. 38 | 39 | ## Preview the Changes 40 | 41 | View the docs web page by opening the HTML in your browser. First, navigate to 42 | the `build/html/` directory and then run the following command: 43 | 44 | ```shell 45 | python -m http.server 46 | ``` 47 | 48 | Afterward, open a web browser and access . 49 | 50 | Check that yours edits formatted correctly and read well. 51 | 52 | ## Special Considerations for Examples and Copydirs 53 | 54 | Freely add examples and images to the `examples` directory. 55 | After you add a notebook or README.md file, update the 56 | `docs/source/toc.yaml` file with entries for the new files. 57 | 58 | Be aware that `README.md` files are renamed to `index.md` 59 | during the build, due to the `copydirs_file_rename` setting. 60 | If you add `examples/blah/README.md`, then add an entry in 61 | the `toc.yaml` file for `examples/blah/index.md`. 62 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/source/_static/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-Merlin/systems/a19d3113e296c69699267bc23f991e5fc4db4ef6/docs/source/_static/.gitkeep -------------------------------------------------------------------------------- /docs/source/_static/NVIDIA-LogoBlack.svg: -------------------------------------------------------------------------------- 1 | NVIDIA-LogoBlack -------------------------------------------------------------------------------- /docs/source/_static/NVIDIA-LogoWhite.svg: -------------------------------------------------------------------------------- 1 | 2 | 13 | 32 | 34 | 36 | 37 | NVIDIA-LogoBlack 39 | 43 | 48 | 53 | 58 | 59 | -------------------------------------------------------------------------------- /docs/source/_static/css/versions.css: -------------------------------------------------------------------------------- 1 | /* Version Switcher */ 2 | 3 | .rst-versions { 4 | flex-align: bottom; 5 | bottom: 0; 6 | left: 0; 7 | z-index: 400 8 | } 9 | 10 | .rst-versions a { 11 | color: var(--nv-green); 12 | text-decoration: none 13 | } 14 | 15 | .rst-versions .rst-badge-small { 16 | display: none 17 | } 18 | 19 | .rst-versions .rst-current-version { 20 | padding: 12px; 21 | display: block; 22 | text-align: right; 23 | font-size: 90%; 24 | cursor: pointer; 25 | border-top: 1px solid rgba(0,0,0,.1); 26 | *zoom:1 27 | } 28 | 29 | .rst-versions .rst-current-version:before,.rst-versions .rst-current-version:after { 30 | display: table; 31 | content: "" 32 | } 33 | 34 | .rst-versions .rst-current-version:after { 35 | clear: both 36 | } 37 | 38 | .rst-versions .rst-current-version .fa-book { 39 | float: left 40 | } 41 | 42 | .rst-versions .rst-current-version .icon-book { 43 | float: left 44 | } 45 | 46 | .rst-versions .rst-current-version.rst-out-of-date { 47 | background-color: #E74C3C; 48 | color: #fff 49 | } 50 | 51 | .rst-versions .rst-current-version.rst-active-old-version { 52 | background-color: #F1C40F; 53 | color: #000 54 | } 55 | 56 | .rst-versions.shift-up { 57 | height: auto; 58 | max-height: 100% 59 | } 60 | 61 | .rst-versions.shift-up .rst-other-versions { 62 | display: block 63 | } 64 | 65 | .rst-versions .rst-other-versions { 66 | font-size: 90%; 67 | padding: 12px; 68 | color: gray; 69 | display: none 70 | } 71 | 72 | .rst-versions .rst-other-versions hr { 73 | display: block; 74 | height: 1px; 75 | border: 0; 76 | margin: 20px 0; 77 | padding: 0; 78 | border-top: solid 1px #413d3d 79 | } 80 | 81 | .rst-versions .rst-other-versions dd { 82 | display: inline-block; 83 | margin: 0 84 | } 85 | 86 | .rst-versions .rst-other-versions dd a { 87 | display: inline-block; 88 | padding: 6px; 89 | color: var(--nv-green); 90 | font-weight: 500; 91 | } 92 | 93 | .rst-versions.rst-badge { 94 | width: auto; 95 | bottom: 20px; 96 | right: 20px; 97 | left: auto; 98 | border: none; 99 | max-width: 300px 100 | } 101 | 102 | .rst-versions.rst-badge .icon-book { 103 | float: none 104 | } 105 | 106 | .rst-versions.rst-badge .fa-book { 107 | float: none 108 | } 109 | 110 | .rst-versions.rst-badge.shift-up .rst-current-version { 111 | text-align: right 112 | } 113 | 114 | .rst-versions.rst-badge.shift-up .rst-current-version .fa-book { 115 | float: left 116 | } 117 | 118 | .rst-versions.rst-badge.shift-up .rst-current-version .icon-book { 119 | float: left 120 | } 121 | 122 | .rst-versions.rst-badge .rst-current-version { 123 | width: auto; 124 | height: 30px; 125 | line-height: 30px; 126 | padding: 0 6px; 127 | display: block; 128 | text-align: center 129 | } 130 | 131 | @media screen and (max-width: 768px) { 132 | .rst-versions { 133 | width:85%; 134 | display: none 135 | } 136 | 137 | .rst-versions.shift { 138 | display: block 139 | } 140 | } 141 | -------------------------------------------------------------------------------- /docs/source/_static/favicon.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-Merlin/systems/a19d3113e296c69699267bc23f991e5fc4db4ef6/docs/source/_static/favicon.png -------------------------------------------------------------------------------- /docs/source/_static/js/rtd-version-switcher.js: -------------------------------------------------------------------------------- 1 | var jQuery = (typeof(window) != 'undefined') ? window.jQuery : require('jquery'); 2 | var doc = $(document); 3 | doc.on('click', "[data-toggle='rst-current-version']", function() { 4 | $("[data-toggle='rst-versions']").toggleClass("shift-up"); 5 | }); 6 | -------------------------------------------------------------------------------- /docs/source/_templates/footer.html: -------------------------------------------------------------------------------- 1 |

2 | Privacy Policy | 3 | Manage My Privacy | 4 | Do Not Sell or Share My Data | 5 | Terms of Service | 6 | Accessibility | 7 | Corporate Policies | 8 | Product Security | 9 | Contact 10 |

11 | -------------------------------------------------------------------------------- /docs/source/_templates/layout.html: -------------------------------------------------------------------------------- 1 | {%- extends "!layout.html" %} 2 | 3 | {%- block extrahead %} 4 | {%- if analytics_id %} 5 | 6 | 7 | 16 | {% endif %} 17 | 18 | 19 | 20 | 21 | {%- endblock %} 22 | -------------------------------------------------------------------------------- /docs/source/_templates/merlin-ecosystem.html: -------------------------------------------------------------------------------- 1 | 14 | -------------------------------------------------------------------------------- /docs/source/_templates/versions.html: -------------------------------------------------------------------------------- 1 | {%- if current_version %} 2 |
3 | 4 | 5 | v: {{ current_version.name }} 6 | 7 | 8 |
9 | {%- if versions.tags %} 10 |
11 |
Tags
12 | {%- for item in versions.tags %} 13 |
{{ item.name }}
14 | {%- endfor %} 15 |
16 | {%- endif %} 17 | {%- if versions.branches %} 18 |
19 |
Branches
20 | {%- for item in versions.branches %} 21 |
{{ item.name }}
22 | {%- endfor %} 23 |
24 | {%- endif %} 25 |
26 |
27 | {%- endif %} 28 | -------------------------------------------------------------------------------- /docs/source/api.rst: -------------------------------------------------------------------------------- 1 | ***************** 2 | API Documentation 3 | ***************** 4 | 5 | .. currentmodule:: merlin.systems 6 | 7 | 8 | Ensemble Graph Constructors 9 | --------------------------- 10 | 11 | .. currentmodule:: merlin.systems.dag 12 | 13 | .. autosummary:: 14 | :toctree: generated 15 | 16 | Ensemble 17 | 18 | Ensemble Operator Constructors 19 | ------------------------------ 20 | 21 | .. currentmodule:: merlin.systems.dag.ops 22 | 23 | .. autosummary:: 24 | :toctree: generated 25 | 26 | workflow.TransformWorkflow 27 | tensorflow.PredictTensorflow 28 | fil.PredictForest 29 | implicit.PredictImplicit 30 | softmax_sampling.SoftmaxSampling 31 | session_filter.FilterCandidates 32 | unroll_features.UnrollFeatures 33 | 34 | .. faiss.QueryFaiss 35 | .. feast.QueryFeast 36 | 37 | 38 | Conversion Functions for Triton Inference Server 39 | ------------------------------------------------ 40 | 41 | .. currentmodule:: merlin.systems.triton 42 | 43 | .. autosummary:: 44 | :toctree: generated 45 | 46 | convert_df_to_triton_input 47 | convert_triton_output_to_df 48 | -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | import os 14 | import re 15 | import subprocess 16 | import sys 17 | from datetime import datetime 18 | 19 | from natsort import natsorted 20 | 21 | sys.path.insert(0, os.path.abspath("../..")) 22 | 23 | repodir = os.path.abspath(os.path.join(__file__, r"../../..")) 24 | gitdir = os.path.join(repodir, r".git") 25 | 26 | 27 | # -- Project information ----------------------------------------------------- 28 | 29 | year_range = "2022" 30 | year_now = str(datetime.now().year) 31 | if year_range != year_now: 32 | year_range = year_range + chr(8211) + year_now 33 | 34 | project = "Merlin Systems" 35 | copyright = year_range + ", NVIDIA" # pylint: disable=W0622 36 | author = "NVIDIA" 37 | 38 | 39 | # -- General configuration --------------------------------------------------- 40 | 41 | # Add any Sphinx extension module names here, as strings. They can be 42 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 43 | # ones. 44 | extensions = [ 45 | "myst_nb", 46 | "sphinx_design", 47 | "sphinx_multiversion", 48 | "sphinx.ext.autodoc", 49 | "sphinx.ext.autosummary", 50 | "sphinx.ext.coverage", 51 | "sphinx.ext.githubpages", 52 | "sphinx.ext.napoleon", 53 | "sphinx.ext.viewcode", 54 | "sphinx.ext.intersphinx", 55 | "sphinx_external_toc", 56 | "sphinxcontrib.copydirs", 57 | ] 58 | 59 | # MyST configuration settings 60 | external_toc_path = "toc.yaml" 61 | myst_enable_extensions = [ 62 | "deflist", 63 | "html_image", 64 | "linkify", 65 | "replacements", 66 | "tasklist", 67 | ] 68 | myst_linkify_fuzzy_links = False 69 | myst_heading_anchors = 3 70 | jupyter_execute_notebooks = "off" 71 | 72 | # Add any paths that contain templates here, relative to this directory. 73 | templates_path = ["_templates"] 74 | 75 | # List of patterns, relative to source directory, that match files and 76 | # directories to ignore when looking for source files. 77 | # This pattern also affects html_static_path and html_extra_path. 78 | exclude_patterns = [] 79 | 80 | # The API documents are RST and include `.. toctree::` directives. 81 | suppress_warnings = ["etoc.toctree"] 82 | 83 | # -- Options for HTML output ------------------------------------------------- 84 | 85 | # The theme to use for HTML and HTML Help pages. See the documentation for 86 | # a list of builtin themes. 87 | # 88 | html_theme = "sphinx_book_theme" 89 | html_title = "Merlin Systems" 90 | html_theme_options = { 91 | "repository_url": "https://github.com/NVIDIA-Merlin/systems", 92 | "use_repository_button": True, 93 | "footer_content_items": ["copyright.html", "footer.html"], 94 | "logo": {"text": "NVIDIA Merlin Systems", "alt_text": "NVIDIA Merlin Systems"}, 95 | } 96 | html_sidebars = { 97 | "**": [ 98 | "navbar-logo.html", 99 | "search-field.html", 100 | "icon-links.html", 101 | "sbt-sidebar-nav.html", 102 | "merlin-ecosystem.html", 103 | "versions.html", 104 | ] 105 | } 106 | html_favicon = "_static/favicon.png" 107 | html_copy_source = True 108 | html_show_sourcelink = False 109 | html_show_sphinx = False 110 | 111 | # Add any paths that contain custom static files (such as style sheets) here, 112 | # relative to this directory. They are copied after the builtin static files, 113 | # so a file named "default.css" will overwrite the builtin "default.css". 114 | html_static_path = ["_static"] 115 | html_css_files = ["css/custom.css", "css/versions.css"] 116 | html_js_files = ["js/rtd-version-switcher.js"] 117 | html_context = {"analytics_id": "G-NVJ1Y1YJHK"} 118 | 119 | source_suffix = [".rst", ".md"] 120 | 121 | if os.path.exists(gitdir): 122 | tag_refs = subprocess.check_output(["git", "tag", "-l", "v*"]).decode("utf-8").split() 123 | tag_refs = [tag for tag in tag_refs if re.match(r"^v[0-9]+.[0-9]+.[0-9]+$", tag)] 124 | tag_refs = natsorted(tag_refs)[-6:] 125 | smv_tag_whitelist = r"^(" + r"|".join(tag_refs) + r")$" 126 | else: 127 | smv_tag_whitelist = r"^v.*$" 128 | 129 | smv_branch_whitelist = r"^(main|stable)$" 130 | 131 | smv_refs_override_suffix = r"-docs" 132 | 133 | intersphinx_mapping = { 134 | "python": ("https://docs.python.org/3", None), 135 | "cudf": ("https://docs.rapids.ai/api/cudf/stable/", None), 136 | "distributed": ("https://distributed.dask.org/en/latest/", None), 137 | "torch": ("https://pytorch.org/docs/stable/", None), 138 | "merlin-core": ("https://nvidia-merlin.github.io/core/stable/", None), 139 | } 140 | 141 | html_baseurl = "https://nvidia-merlin.github.io/systems/stable/" 142 | 143 | autodoc_inherit_docstrings = False 144 | autodoc_default_options = { 145 | "members": True, 146 | "undoc-members": True, 147 | "show-inheritance": False, 148 | "member-order": "bysource", 149 | } 150 | 151 | autosummary_generate = True 152 | 153 | copydirs_additional_dirs = [ 154 | "../../LICENSE", 155 | "../../README.md", 156 | "../../examples/", 157 | ] 158 | copydirs_file_rename = {"README.md": "index.md"} 159 | -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- 1 | Merlin Systems 2 | ============== 3 | 4 | The Merlin Systems library provides facilities that simplify building a recommender system for batch or serving. 5 | 6 | To learn more, start with the `Introduction `_. 7 | 8 | Related Resources 9 | ----------------- 10 | 11 | Merlin Systems GitHub Repository 12 | ``_ 13 | 14 | About Merlin 15 | Merlin is the overarching project that brings together the Merlin projects. 16 | See the `documentation `_ 17 | or the `repository `_ on GitHub. 18 | 19 | Developer website for Merlin 20 | More information about Merlin is available at our developer website: 21 | ``_. 22 | 23 | Index 24 | ----- 25 | 26 | * :ref:`genindex` 27 | -------------------------------------------------------------------------------- /docs/source/toc.yaml: -------------------------------------------------------------------------------- 1 | root: index 2 | subtrees: 3 | - caption: Contents 4 | entries: 5 | - file: README.md 6 | title: Introduction 7 | - file: examples/index 8 | title: Example Notebooks 9 | subtrees: 10 | - titlesonly: True 11 | entries: 12 | - file: examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb 13 | - file: examples/Serving-An-XGboost-Model-With-Merlin-Systems.ipynb 14 | - file: examples/Serving-An-Implicit-Model-With-Merlin-Systems.ipynb 15 | - file: api 16 | title: API Documentation 17 | -------------------------------------------------------------------------------- /merlin/systems/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import warnings 17 | 18 | from . import _version 19 | 20 | # suppress some warnings with cudf warning about column ordering with dlpack 21 | # and numba warning about deprecated environment variables 22 | warnings.filterwarnings("ignore", module="cudf.io.dlpack") 23 | warnings.filterwarnings("ignore", module="numba.cuda.envvars") 24 | 25 | # cudf warns about column ordering with dlpack methods, ignore it 26 | warnings.filterwarnings("ignore", module="cudf.io.dlpack") 27 | 28 | 29 | __version__ = _version.get_versions()["version"] 30 | -------------------------------------------------------------------------------- /merlin/systems/dag/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | # alias submodules here to avoid breaking everything with moving to submodules 18 | # flake8: noqa 19 | from .ensemble import Ensemble 20 | -------------------------------------------------------------------------------- /merlin/systems/dag/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | 18 | def compute_dims(col_schema): 19 | """ 20 | Compute Triton dimensions for a column from its schema 21 | 22 | Parameters 23 | ---------- 24 | col_schema : ColumnSchema 25 | Schema of the column to compute dimensions for 26 | 27 | Returns 28 | ------- 29 | List[int] 30 | Triton dimensions for the column 31 | """ 32 | dims = [-1] 33 | 34 | if col_schema.shape is not None and col_schema.shape.dims is not None: 35 | for dim in col_schema.shape.as_tuple[1:]: 36 | dim = dim if isinstance(dim, int) else -1 37 | dims.append(dim) 38 | 39 | return dims 40 | -------------------------------------------------------------------------------- /merlin/systems/dag/ops/compat.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # pylint: disable=unused-import 17 | 18 | try: 19 | import xgboost 20 | except ImportError: 21 | xgboost = None 22 | try: 23 | import sklearn.ensemble as sklearn_ensemble 24 | except ImportError: 25 | sklearn_ensemble = None 26 | try: 27 | import treelite.sklearn as treelite_sklearn 28 | except ImportError: 29 | treelite_sklearn = None 30 | try: 31 | import treelite.Model as treelite_model 32 | except ImportError: 33 | treelite_model = None 34 | try: 35 | import lightgbm 36 | except ImportError: 37 | lightgbm = None 38 | try: 39 | import cuml.ensemble as cuml_ensemble 40 | except ImportError: 41 | cuml_ensemble = None 42 | try: 43 | from cuml import ForestInference as cuml_fil 44 | except ImportError: 45 | cuml_fil = None 46 | try: 47 | import triton_python_backend_utils as pb_utils 48 | except ImportError: 49 | pb_utils = None 50 | -------------------------------------------------------------------------------- /merlin/systems/dag/ops/implicit.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import importlib 17 | import pathlib 18 | 19 | import numpy as np 20 | 21 | from merlin.core.protocols import Transformable 22 | from merlin.dag import BaseOperator, ColumnSelector 23 | from merlin.schema import ColumnSchema, Schema 24 | 25 | try: 26 | import implicit 27 | from packaging.version import Version 28 | 29 | if Version(implicit.__version__) < Version("0.6.0"): 30 | raise RuntimeError( 31 | "Implicit version 0.6.0 or higher required. (for model save/load methods)." 32 | ) 33 | except ImportError: 34 | implicit = None 35 | 36 | 37 | class PredictImplicit(BaseOperator): 38 | """Operator for running inference on Implicit models..""" 39 | 40 | def __init__(self, model, num_to_recommend: int = 10, **kwargs): 41 | """Instantiate an Implicit prediction operator. 42 | 43 | Parameters 44 | ---------- 45 | model : An Implicit Model instance 46 | num_to_recommend : int 47 | the number of items to return 48 | """ 49 | self.model = model 50 | self.model_module_name: str = self.model.__module__ 51 | self.model_class_name: str = self.model.__class__.__name__ 52 | self.num_to_recommend = num_to_recommend 53 | super().__init__(**kwargs) 54 | 55 | def __getstate__(self): 56 | return {k: v for k, v in self.__dict__.items() if k != "model"} 57 | 58 | def load_artifacts(self, artifact_path: str): 59 | if artifact_path: 60 | model_file = pathlib.Path(artifact_path) / "model.npz" 61 | 62 | model_module_name = self.model_module_name 63 | model_class_name = self.model_class_name 64 | model_module = importlib.import_module(model_module_name) 65 | model_cls = getattr(model_module, model_class_name) 66 | 67 | self.model = model_cls.load(str(model_file)) 68 | 69 | def save_artifacts(self, artifact_path: str): 70 | model_path = pathlib.Path(artifact_path) / "model.npz" 71 | self.model.save(str(model_path)) 72 | 73 | def compute_input_schema( 74 | self, 75 | root_schema: Schema, 76 | parents_schema: Schema, 77 | deps_schema: Schema, 78 | selector: ColumnSelector, 79 | ) -> Schema: 80 | """Return the input schema representing the input columns this operator expects to use.""" 81 | return Schema([ColumnSchema("user_id", dtype="int64")]) 82 | 83 | def compute_output_schema( 84 | self, 85 | input_schema: Schema, 86 | col_selector: ColumnSelector, 87 | prev_output_schema: Schema = None, 88 | ) -> Schema: 89 | """Return the output schema representing the columns this operator returns.""" 90 | return Schema([ColumnSchema("ids", dtype="int64"), ColumnSchema("scores", dtype="float64")]) 91 | 92 | def transform( 93 | self, col_selector: ColumnSelector, transformable: Transformable 94 | ) -> Transformable: 95 | """Transform the dataframe by applying this operator to the set of input columns. 96 | 97 | Parameters 98 | ----------- 99 | df: TensorTable 100 | A pandas or cudf dataframe that this operator will work on 101 | 102 | Returns 103 | ------- 104 | TensorTable 105 | Returns a transformed dataframe for this operator""" 106 | user_id = transformable["user_id"].values.ravel() 107 | user_items = None 108 | ids, scores = self.model.recommend( 109 | user_id, user_items, N=self.num_to_recommend, filter_already_liked_items=False 110 | ) 111 | return type(transformable)( 112 | {"ids": ids.astype(np.int64), "scores": scores.astype(np.float64)} 113 | ) 114 | -------------------------------------------------------------------------------- /merlin/systems/dag/ops/pytorch.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import os 17 | 18 | import numpy as np 19 | 20 | from merlin.core.compat.torch import torch 21 | from merlin.core.protocols import Transformable 22 | from merlin.dag import BaseOperator, ColumnSelector 23 | from merlin.schema import Schema 24 | from merlin.table import TensorTable 25 | 26 | 27 | class PredictPyTorch(BaseOperator): 28 | """ 29 | This operator takes a pytorch model and packages it correctly for tritonserver 30 | to run, on the pytorch backend. 31 | """ 32 | 33 | def __init__(self, model_or_path, input_schema: Schema, output_schema: Schema, backend="torch"): 34 | """ 35 | Instantiate a PredictPyTorch inference operator. 36 | 37 | Parameters 38 | ---------- 39 | model_or_path : PyTorch model or string 40 | This can be a pytorch model or a path to a pytorch model. 41 | input_schema : Schema 42 | Input schema for the pytorch model. This could be the output schema of the NVTabular 43 | workflow that produced your training data. 44 | output_schema : Schema 45 | Output schema for the pytorch model. 46 | """ 47 | 48 | super().__init__() 49 | self._torch_model_name = None 50 | self.input_schema = input_schema 51 | self.output_schema = output_schema 52 | 53 | if model_or_path is not None: 54 | if isinstance(model_or_path, (str, os.PathLike)): 55 | self.path = model_or_path 56 | self.model = torch.load(self.path) 57 | else: 58 | self.path = None 59 | self.model = model_or_path 60 | 61 | # This is a hack to enable the ensemble to use the same shape as this 62 | # these lines mutate the input and output schema with the additional property 63 | # `triton_scalar_shape` which represents the expected shape for this feature 64 | for col_name, col_schema in self.input_schema.column_schemas.items(): 65 | self.input_schema[col_name] = col_schema.with_properties( 66 | {"triton_scalar_shape": self.scalar_shape} 67 | ) 68 | 69 | for col_name, col_schema in self.output_schema.column_schemas.items(): 70 | self.output_schema[col_name] = col_schema.with_properties( 71 | {"triton_scalar_shape": self.scalar_shape} 72 | ) 73 | 74 | def __getstate__(self): 75 | return {k: v for k, v in self.__dict__.items() if k != "model"} 76 | 77 | def compute_input_schema( 78 | self, 79 | root_schema: Schema, 80 | parents_schema: Schema, 81 | deps_schema: Schema, 82 | selector: ColumnSelector, 83 | ) -> Schema: 84 | """ 85 | Use the input schema supplied during object creation. 86 | """ 87 | return self.input_schema 88 | 89 | def compute_output_schema( 90 | self, input_schema: Schema, col_selector: ColumnSelector, prev_output_schema: Schema = None 91 | ) -> Schema: 92 | """ 93 | Use the output schema supplied during object creation. 94 | """ 95 | return self.output_schema 96 | 97 | def transform(self, col_selector: ColumnSelector, transformable: Transformable): 98 | output_type = type(transformable) 99 | if not isinstance(transformable, TensorTable): 100 | transformable = TensorTable.from_df(transformable) 101 | 102 | tensor_dict = {} 103 | for column in transformable.columns: 104 | tensor_dict[column] = torch.from_numpy(np.squeeze(transformable[column].values)) 105 | 106 | result = self.model(tensor_dict) 107 | output = {} 108 | for idx, col in enumerate(self.output_schema.column_names): 109 | output[col] = result[:, idx].detach().numpy() 110 | 111 | output = TensorTable(output) 112 | if not isinstance(output, output_type): 113 | output = output.to_df() 114 | 115 | return output 116 | 117 | @property 118 | def scalar_shape(self): 119 | return [] 120 | -------------------------------------------------------------------------------- /merlin/systems/dag/ops/softmax_sampling.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from merlin.core.protocols import Transformable 4 | from merlin.dag import BaseOperator 5 | from merlin.dag.node import Node 6 | from merlin.dag.selector import ColumnSelector 7 | from merlin.schema import ColumnSchema, Schema 8 | 9 | 10 | class SoftmaxSampling(BaseOperator): 11 | """ 12 | Given inputs of ID and prediction, this operator will sort all 13 | inputs in descending order. 14 | """ 15 | 16 | def __init__(self, relevance_col, temperature=20.0, topk=10, _input_col=None): 17 | """ 18 | Create a SoftmaxSampling Pipelineable Inference Operator. 19 | 20 | Parameters 21 | ---------- 22 | relevance_col : string 23 | The column to judge sorting order with. 24 | temperature : float, optional 25 | Value which will be used to effect the weights used in sorting, by default 20.0 26 | topk : int, optional 27 | The max number of results you wish to receive as output, by default 10 28 | _input_col : _type_, optional 29 | The column whose values will be sorted, by default None. 30 | """ 31 | self.relevance_col = Node.construct_from(relevance_col) 32 | self.temperature = temperature 33 | self.topk = topk 34 | self._input_col_name = _input_col 35 | self._relevance_col_name = relevance_col 36 | super().__init__() 37 | 38 | @property 39 | def dependencies(self): 40 | return self.relevance_col 41 | 42 | def compute_input_schema( 43 | self, 44 | root_schema: Schema, 45 | parents_schema: Schema, 46 | deps_schema: Schema, 47 | selector: ColumnSelector, 48 | ) -> Schema: 49 | input_schema = super().compute_input_schema( 50 | root_schema, parents_schema, deps_schema, selector 51 | ) 52 | if len(parents_schema.column_schemas) > 1: 53 | raise ValueError( 54 | "More than one input has been detected for this node," 55 | f" inputs received: {input_schema.column_names}" 56 | ) 57 | 58 | self._input_col_name = parents_schema.column_names[0] 59 | self._relevance_col_name = deps_schema.column_names[0] 60 | return input_schema 61 | 62 | def compute_output_schema( 63 | self, input_schema: Schema, col_selector: ColumnSelector, prev_output_schema: Schema = None 64 | ) -> Schema: 65 | """Describe the operator's outputs""" 66 | return Schema( 67 | [ 68 | ColumnSchema( 69 | "ordered_ids", dtype=input_schema[self._input_col_name].dtype, dims=(None, 1) 70 | ), 71 | ColumnSchema( 72 | "ordered_scores", 73 | dtype=input_schema[self._relevance_col_name].dtype, 74 | dims=(None, 1), 75 | ), 76 | ] 77 | ) 78 | 79 | def transform( 80 | self, col_selector: ColumnSelector, transformable: Transformable 81 | ) -> Transformable: 82 | """Transform the dataframe by applying this operator to the set of input columns""" 83 | # Extract parameters from the request 84 | candidate_ids = transformable[self._input_col_name].values.reshape(-1) 85 | predicted_scores = transformable[self._relevance_col_name].values.reshape(-1) 86 | 87 | # Exponential sort trick for sampling from a distribution without replacement from: 88 | 89 | # Pavlos S. Efraimidis, Paul G. Spirakis, Weighted random sampling with a reservoir, 90 | # Information Processing Letters, Volume 97, Issue 5, 2006, Pages 181-185, ISSN 0020-0190, 91 | # https://doi.org/10.1016/j.ipl.2005.11.003. 92 | 93 | # As implemented by Tim Vieira in "Algorithms for sampling without replacement" 94 | # https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/ 95 | 96 | # The weights for the sampling distribution are the softmax of the scores 97 | weights = np.exp(self.temperature * predicted_scores) / np.sum(predicted_scores) 98 | 99 | # This is the core of the exponential sampling trick, which creates a 100 | # set of values that depend on both the predicted scores and random 101 | # variables, resulting in a set of values that will sort into an order 102 | # that reflects sampling without replacement according to the weight 103 | # distribution 104 | num_items = candidate_ids.shape[0] 105 | exponentials = -np.log(np.random.uniform(0, 1, size=(num_items,))) 106 | exponentials /= weights 107 | 108 | # This is just bookkeeping to produce the final ordered list of recs 109 | sorted_indices = np.argsort(exponentials) 110 | topk_item_ids = candidate_ids[sorted_indices][: self.topk] 111 | topk_item_scores = predicted_scores[sorted_indices][: self.topk] 112 | ordered_item_ids = topk_item_ids.reshape(1, -1) 113 | ordered_item_scores = topk_item_scores.reshape(1, -1) 114 | 115 | return type(transformable)( 116 | {"ordered_ids": ordered_item_ids, "ordered_scores": ordered_item_scores} 117 | ) 118 | -------------------------------------------------------------------------------- /merlin/systems/dag/ops/unroll_features.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import numpy as np 18 | 19 | from merlin.core.protocols import Transformable 20 | from merlin.dag import BaseOperator, DataFormats, Node 21 | from merlin.dag.selector import ColumnSelector 22 | from merlin.schema import Schema 23 | 24 | 25 | class UnrollFeatures(BaseOperator): 26 | """ 27 | This operator takes a target column and joins the "unroll" columns to the target. This helps 28 | when broadcasting a series of user features to a set of items. 29 | """ 30 | 31 | def __init__(self, item_id_col, unroll_cols, unrolled_prefix=""): 32 | self.item_id_col = item_id_col 33 | self.unroll_cols = Node.construct_from(unroll_cols) 34 | self.unrolled_prefix = unrolled_prefix 35 | super().__init__() 36 | 37 | @property 38 | def dependencies(self): 39 | return self.unroll_cols 40 | 41 | def compute_output_schema( 42 | self, input_schema: Schema, col_selector: ColumnSelector, prev_output_schema: Schema = None 43 | ) -> Schema: 44 | schema = super().compute_output_schema(input_schema, col_selector, prev_output_schema) 45 | 46 | for col_name, col_schema in self.unroll_cols.output_schema.column_schemas.items(): 47 | schema.column_schemas.pop(col_name, None) 48 | col_name = f"{self.unrolled_prefix}_{col_name}" if self.unrolled_prefix else col_name 49 | schema[col_name] = col_schema.with_name(col_name) 50 | 51 | return schema 52 | 53 | def transform( 54 | self, col_selector: ColumnSelector, transformable: Transformable 55 | ) -> Transformable: 56 | num_items = transformable[self.item_id_col].shape[0].max 57 | outputs = {} 58 | for col_name, col_value in transformable.items(): 59 | outputs[col_name] = col_value 60 | 61 | for col in self._unroll_col_names: 62 | target = outputs.pop(col) 63 | col_name = f"{self.unrolled_prefix}_{col}" if self.unrolled_prefix else col 64 | outputs[col_name] = np.repeat(target, num_items, axis=0) 65 | 66 | return type(transformable)(outputs) 67 | 68 | @property 69 | def _unroll_col_names(self): 70 | if self.unroll_cols.selector: 71 | return self.unroll_cols.selector.names 72 | else: 73 | return self.unroll_cols.output_columns.names 74 | 75 | @property 76 | def supported_formats(self) -> DataFormats: 77 | return DataFormats.NUMPY_TENSOR_TABLE | DataFormats.CUPY_TENSOR_TABLE 78 | -------------------------------------------------------------------------------- /merlin/systems/dag/ops/workflow.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | from typing import List 17 | 18 | from merlin.core.protocols import Transformable 19 | from merlin.dag import BaseOperator, ColumnSelector, DataFormats 20 | from merlin.schema import Schema 21 | from merlin.table import TensorTable 22 | 23 | 24 | class TransformWorkflow(BaseOperator): 25 | """ 26 | This operator takes a workflow and turns it into a ensemble operator so that we can 27 | execute feature engineering during ensemble on tritonserver. 28 | """ 29 | 30 | def __init__( 31 | self, 32 | workflow=None, 33 | sparse_max: dict = None, 34 | max_batch_size: int = None, 35 | label_columns: List[str] = None, 36 | cats: List[str] = None, 37 | conts: List[str] = None, 38 | ): 39 | """ 40 | Creates a Transform Workflow operator for a target workflow. 41 | 42 | Parameters 43 | ---------- 44 | workflow : Nvtabular.Workflow 45 | The workflow to transform data in ensemble. 46 | sparse_max : dict, optional 47 | Dictionary representing key(name)/val(max value) pairs of max sparsity, by default None 48 | max_batch_size : int, optional 49 | Maximum batch size, by default None 50 | label_columns : List[str], optional 51 | List of strings identifying the label columns, by default None 52 | cats : List[str], optional 53 | List of strings identifying categorical columns, by default None 54 | conts : List[str], optional 55 | List of string identifying continuous columns, by default None 56 | """ 57 | super().__init__() 58 | 59 | self.workflow = workflow 60 | if label_columns: 61 | self.workflow = workflow.remove_inputs(label_columns) 62 | self._nvt_model_name = None 63 | self.sparse_max = sparse_max or {} 64 | self.max_batch_size = max_batch_size 65 | self.label_columns = label_columns or [] 66 | self.cats = cats or [] 67 | self.conts = conts or [] 68 | 69 | if self.workflow is not None: 70 | self.input_schema = self.workflow.input_schema 71 | self.output_schema = self.workflow.output_schema 72 | 73 | @property 74 | def nvt_model_name(self): 75 | return self._nvt_model_name 76 | 77 | def set_nvt_model_name(self, nvt_model_name): 78 | self._nvt_model_name = nvt_model_name 79 | 80 | def compute_output_schema( 81 | self, input_schema: Schema, col_selector: ColumnSelector, prev_output_schema: Schema = None 82 | ) -> Schema: 83 | """Returns output schema of operator""" 84 | return self.workflow.output_schema 85 | 86 | def transform( 87 | self, col_selector: ColumnSelector, transformable: Transformable 88 | ) -> Transformable: 89 | """Run nvtabular workflow transformations. 90 | 91 | Parameters 92 | ---------- 93 | col_selector : ColumnSelector 94 | Unused ColumunSelector input 95 | transformable : Transformable 96 | Input features to model 97 | 98 | Returns 99 | ------- 100 | Transformable 101 | workflow transform 102 | """ 103 | output_type = type(transformable) 104 | if isinstance(transformable, TensorTable): 105 | transformable = transformable.to_df() 106 | 107 | output = self.workflow._transform_df(transformable) 108 | 109 | if not isinstance(output, output_type): 110 | output = TensorTable.from_df(output) 111 | 112 | return output 113 | 114 | @property 115 | def supported_formats(self) -> DataFormats: 116 | return DataFormats.PANDAS_DATAFRAME | DataFormats.CUDF_DATAFRAME 117 | -------------------------------------------------------------------------------- /merlin/systems/dag/runtimes/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2023, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # flake8: noqa 17 | -------------------------------------------------------------------------------- /merlin/systems/dag/runtimes/triton/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # flake8: noqa 17 | from merlin.systems.dag.runtimes.triton.runtime import TritonExecutorRuntime # noqa 18 | -------------------------------------------------------------------------------- /merlin/systems/dag/runtimes/triton/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /merlin/systems/dag/runtimes/triton/ops/operator.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | from abc import ABCMeta, abstractmethod 17 | 18 | import tritonclient.grpc.model_config_pb2 as model_config 19 | 20 | from merlin.core.protocols import Transformable 21 | from merlin.dag import BaseOperator 22 | from merlin.dag.selector import ColumnSelector 23 | from merlin.schema import Schema 24 | from merlin.systems.triton.export import _convert_dtype 25 | 26 | 27 | class TritonOperator(BaseOperator, metaclass=ABCMeta): 28 | """Base class for Triton operators.""" 29 | 30 | def __init__(self, base_op: BaseOperator): 31 | """Construct TritonOperator from a base operator. 32 | 33 | Parameters 34 | ---------- 35 | base_op : merlin.systems.dag.ops.operator.InfereneOperator 36 | Base operator used to construct this Triton op. 37 | """ 38 | super().__init__() 39 | self.op = base_op 40 | 41 | @property 42 | def export_name(self): 43 | """ 44 | Provides a clear common english identifier for this operator. 45 | 46 | Returns 47 | ------- 48 | String 49 | Name of the current class as spelled in module. 50 | """ 51 | return self.__class__.__name__.lower() 52 | 53 | def transform( 54 | self, col_selector: ColumnSelector, transformable: Transformable 55 | ) -> Transformable: 56 | """Transform the dataframe by applying this operator to the set of input columns 57 | 58 | Parameters 59 | ----------- 60 | df: Dataframe 61 | A pandas or cudf dataframe that this operator will work on 62 | 63 | Returns 64 | ------- 65 | DataFrame 66 | Returns a transformed dataframe for this operator 67 | """ 68 | return transformable 69 | 70 | @abstractmethod 71 | def export( 72 | self, 73 | path: str, 74 | input_schema: Schema, 75 | output_schema: Schema, 76 | params: dict = None, 77 | node_id: int = None, 78 | version: int = 1, 79 | ): 80 | """ 81 | Export the Operator to as a Triton Model at the path corresponding to the model repository. 82 | 83 | Parameters 84 | ---------- 85 | path : str 86 | Artifact export path 87 | input_schema : Schema 88 | A schema with information about the inputs to this operator. 89 | output_schema : Schema 90 | A schema with information about the outputs of this operator. 91 | params : dict, optional 92 | Parameters dictionary of key, value pairs stored in exported config, by default None. 93 | node_id : int, optional 94 | The placement of the node in the graph (starts at 1), by default None. 95 | version : int, optional 96 | The version of the model, by default 1. 97 | 98 | Returns 99 | ------- 100 | model_config: ModelConfig 101 | The config for the operator (model) if defined. 102 | """ 103 | 104 | 105 | def add_model_param(params, paramclass, col_schema, dims=None): 106 | if col_schema.is_list and col_schema.is_ragged: 107 | params.append( 108 | paramclass( 109 | name=col_schema.name + "__values", 110 | data_type=_convert_dtype(col_schema.dtype), 111 | dims=dims[1:], 112 | ) 113 | ) 114 | params.append( 115 | paramclass( 116 | name=col_schema.name + "__offsets", data_type=model_config.TYPE_INT32, dims=[-1] 117 | ) 118 | ) 119 | else: 120 | params.append( 121 | paramclass(name=col_schema.name, data_type=_convert_dtype(col_schema.dtype), dims=dims) 122 | ) 123 | -------------------------------------------------------------------------------- /merlin/systems/model_registry.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import requests 4 | 5 | 6 | class ModelRegistry(ABC): 7 | """ 8 | The ModelRegistry class is used to find model paths that will be imported into an 9 | InferenceOperator. 10 | 11 | To implement your own ModelRegistry subclass, the only method that must be implemented is 12 | `get_artifact_uri`, which must return a string indicating the model's export path. 13 | 14 | ```python 15 | PredictTensorflow.from_model_registry( 16 | MyModelRegistry("model_name", "model_version") 17 | ) 18 | ``` 19 | """ 20 | 21 | @abstractmethod 22 | def get_artifact_uri(self) -> str: 23 | """ 24 | This returns the URI of the model artifact. 25 | """ 26 | 27 | 28 | class MLFlowModelRegistry(ModelRegistry): 29 | def __init__(self, name: str, version: str, tracking_uri: str): 30 | """ 31 | Fetches the model path from an mlflow model registry. 32 | 33 | Note that this will return a relative path if you did not configure your mlflow 34 | experiment's `artifact_location` to be an absolute path. 35 | 36 | Parameters 37 | ---------- 38 | name : str 39 | Name of the model in the mlflow registry. 40 | version : str 41 | Version of the model to use. 42 | tracking_uri : str 43 | Base URI of the mlflow tracking server. If running locally, this would likely be 44 | http://localhost:5000 45 | """ 46 | self.name = name 47 | self.version = version 48 | self.tracking_uri = tracking_uri.rstrip("/") 49 | 50 | def get_artifact_uri(self) -> str: 51 | mv = requests.get( 52 | f"{self.tracking_uri}/ajax-api/2.0/preview/mlflow/model-versions/get-download-uri", 53 | params={"name": self.name, "version": self.version}, 54 | ) 55 | 56 | if mv.status_code != 200: 57 | raise ValueError( 58 | f"Could not find a Model Version for model {self.name} with version {self.version}." 59 | ) 60 | model_path = mv.json()["artifact_uri"] 61 | return model_path 62 | -------------------------------------------------------------------------------- /merlin/systems/triton/export.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # 15 | import os 16 | 17 | # this needs to be before any modules that import protobuf 18 | os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" 19 | 20 | import tritonclient.grpc.model_config_pb2 as model_config # noqa 21 | 22 | import merlin.dtypes as md # noqa 23 | from merlin.core.dispatch import is_string_dtype # noqa 24 | from merlin.systems.dag.ops import compute_dims # noqa 25 | 26 | 27 | def _add_model_param(col_schema, paramclass, params, dims=None): 28 | dims = dims if dims is not None else compute_dims(col_schema) 29 | if col_schema.is_list and col_schema.is_ragged: 30 | params.append( 31 | paramclass( 32 | name=col_schema.name + "__values", 33 | data_type=_convert_dtype(col_schema.dtype), 34 | dims=[-1], 35 | ) 36 | ) 37 | params.append( 38 | paramclass( 39 | name=col_schema.name + "__offsets", data_type=model_config.TYPE_INT32, dims=[-1] 40 | ) 41 | ) 42 | else: 43 | params.append( 44 | paramclass(name=col_schema.name, data_type=_convert_dtype(col_schema.dtype), dims=dims) 45 | ) 46 | 47 | 48 | def _convert_dtype(dtype): 49 | """converts a dtype to the appropriate triton proto type""" 50 | dtype = md.dtype(dtype) 51 | try: 52 | return dtype.to("triton") 53 | except ValueError: 54 | dtype = dtype.to_numpy 55 | 56 | if is_string_dtype(dtype): 57 | return model_config.TYPE_STRING 58 | else: 59 | raise ValueError(f"Can't convert {dtype} to a Triton dtype") 60 | 61 | 62 | def _convert_string2pytorch_dtype(dtype): 63 | """converts a dtype to the appropriate torch type""" 64 | 65 | import torch 66 | 67 | if not isinstance(dtype, str): 68 | dtype_name = dtype.name 69 | else: 70 | dtype_name = dtype 71 | 72 | dtypes = { 73 | "TYPE_FP64": torch.float64, 74 | "TYPE_FP32": torch.float32, 75 | "TYPE_FP16": torch.float16, 76 | "TYPE_INT64": torch.int64, 77 | "TYPE_INT32": torch.int32, 78 | "TYPE_INT16": torch.int16, 79 | "TYPE_INT8": torch.int8, 80 | "TYPE_UINT8": torch.uint8, 81 | "TYPE_BOOL": torch.bool, 82 | } 83 | 84 | if is_string_dtype(dtype): 85 | return model_config.TYPE_STRING 86 | elif dtype_name in dtypes: 87 | return dtypes[dtype_name] 88 | else: 89 | raise ValueError(f"Can't convert dtype {dtype})") 90 | -------------------------------------------------------------------------------- /merlin/systems/triton/models/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /merlin/systems/workflow/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # 15 | from merlin.schema import Tags 16 | 17 | 18 | def get_embedding_sizes(source, output_dtypes=None): 19 | """Returns a dictionary of embedding sizes from a workflow or workflow_node 20 | 21 | Parameters 22 | ---------- 23 | source : Workflow or ColumnSelector 24 | Either a nvtabular Workflow or ColumnSelector object that we should use to find 25 | embedding sizes 26 | output_dtypes : dict, optional 27 | Optional dictionary of column_name:dtype. If passing a workflow object dtypes 28 | will be read from the workflow. This is used to figure out which columns 29 | are multihot-categorical, which are split out by this function. If passed a workflow_node 30 | and this parameter isn't set, you won't have multihot columns returned separately 31 | """ 32 | # TODO: do we need to distinguish multihot columns here? (if so why? ) 33 | 34 | # have to lazy import Workflow to avoid circular import errors 35 | from nvtabular.workflow import Workflow 36 | 37 | output_node = source.output_node if isinstance(source, Workflow) else source 38 | 39 | if isinstance(source, Workflow): 40 | output_dtypes = output_dtypes or source.output_dtypes 41 | else: 42 | # passed in a column group 43 | output_dtypes = output_dtypes or {} 44 | 45 | output = {} 46 | multihot_columns = set() 47 | cats_schema = output_node.output_schema.select_by_tag(Tags.CATEGORICAL) 48 | for col_name, col_schema in cats_schema.column_schemas.items(): 49 | if col_schema.dtype and col_schema.is_list and col_schema.is_ragged: 50 | # multi hot so remove from output and add to multihot 51 | multihot_columns.add(col_name) 52 | 53 | embeddings_sizes = col_schema.properties.get("embedding_sizes", {}) 54 | cardinality = embeddings_sizes["cardinality"] 55 | dimensions = embeddings_sizes["dimension"] 56 | output[col_name] = (cardinality, dimensions) 57 | 58 | # TODO: returning different return types like this (based off the presence 59 | # of multihot features) is pretty janky. fix. 60 | if not multihot_columns: 61 | return output 62 | 63 | single_hots = {k: v for k, v in output.items() if k not in multihot_columns} 64 | multi_hots = {k: v for k, v in output.items() if k in multihot_columns} 65 | return single_hots, multi_hots 66 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = [ 3 | "setuptools>=42", 4 | "wheel", 5 | ] 6 | build-backend = "setuptools.build_meta" 7 | 8 | [tool.black] 9 | line-length = 100 10 | 11 | [tool.isort] 12 | use_parentheses = true 13 | multi_line_output = 3 14 | include_trailing_comma = true 15 | force_grid_wrap = 0 16 | ensure_newline_before_comments = true 17 | line_length = 100 18 | balanced_wrapping = true 19 | indent = " " 20 | known_third_party = ["cudf", "cupy", "dask", "dask_cuda", "dask_cudf", "numba", "numpy", "pytest", "torch", "rmm", "tensorflow", "xgboost", "lightgbm", "sklearn"] 21 | known_first_party = ["merlin", "nvtabular"] 22 | skip = ["build", ".eggs"] 23 | 24 | [tool.interrogate] 25 | ignore-init-method = true 26 | ignore-init-module = true 27 | ignore-magic = true 28 | ignore-module = true 29 | ignore-private = true 30 | ignore-property-decorators = true 31 | ignore-nested-classes = true 32 | ignore-nested-functions = true 33 | ignore-semiprivate = true 34 | ignore-setters = true 35 | fail-under = 70 36 | exclude = ["build", "docs", "examples", "tests", "setup.py", "versioneer.py"] 37 | verbose = 1 38 | omit-covered-files = true 39 | quiet = false 40 | whitelist-regex = [] 41 | ignore-regex = [] 42 | color = true 43 | 44 | [tool.pytest.ini_options] 45 | filterwarnings = [ 46 | 'ignore:`np.*` is a deprecated alias:DeprecationWarning', 47 | 'ignore:The default dtype for empty Series:DeprecationWarning', 48 | 'ignore:General-metadata information not detected:UserWarning', 49 | 'ignore:Changing an NVTabular Dataset to CPU mode:UserWarning', 50 | 'ignore:Initializing an NVTabular Dataset in CPU mode:UserWarning', 51 | 'ignore:Performing a hash-based transformation:UserWarning', 52 | 'ignore:WARNING..cuDF.to_dlpack', 53 | 'ignore:::numba.cuda.envvar:', 54 | 'ignore:Call to deprecated create function:DeprecationWarning', 55 | 'ignore:Only created .* files did not have enough partitions to create .* file:UserWarning', 56 | 'ignore:distutils Version classes are deprecated. Use packaging.version instead:DeprecationWarning', 57 | ] 58 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | markers = 3 | notebook: mark as testing notebooks 4 | version: mark as testing the version number 5 | premerge: mark as running only before a PR is merged 6 | postmerge: mark as running only after a PR is merged 7 | 8 | -------------------------------------------------------------------------------- /requirements/base.txt: -------------------------------------------------------------------------------- 1 | merlin-core>=23.4.0 2 | nvtabular>=23.4.0 3 | requests>=2.10,<3 4 | # regardless of whether in GPU/CPU environment must support sklearn which is CPU only. 5 | treelite==2.4.0 6 | treelite_runtime==2.4.0 -------------------------------------------------------------------------------- /requirements/dev.txt: -------------------------------------------------------------------------------- 1 | # Required to generate docs and also run tests 2 | tritonclient[all] 3 | -------------------------------------------------------------------------------- /requirements/docs.txt: -------------------------------------------------------------------------------- 1 | -r base.txt 2 | -r dev.txt 3 | 4 | tritonclient[all] 5 | tensorflow<=2.9.0 6 | 7 | sphinx-book-theme~=1.0.1 8 | sphinx-multiversion@git+https://github.com/mikemckiernan/sphinx-multiversion.git 9 | sphinxcontrib-copydirs@git+https://github.com/mikemckiernan/sphinxcontrib-copydirs.git 10 | recommonmark~=0.7.1 11 | Jinja2<3.1 12 | natsort~=8.4.0 13 | myst-nb~=0.17.2 14 | linkify-it-py~=2.0.3 15 | sphinx-external-toc~=0.3.1 16 | attrs~=23.2.0 17 | sphinx_design~=0.5.0 18 | 19 | # keep support for numpy builtin type aliases for previous tags 20 | # numpy builtin aliases like np.str were removed in 1.24 21 | # This can be unpinned when we no longer build docs for versions of Merlin prior 23.05 22 | numpy<1.24 23 | -------------------------------------------------------------------------------- /requirements/gpu.txt: -------------------------------------------------------------------------------- 1 | faiss-gpu 2 | dask>=2022.3.0 3 | distributed>=2022.3.0 4 | -------------------------------------------------------------------------------- /requirements/test-cpu.txt: -------------------------------------------------------------------------------- 1 | -r test.txt 2 | 3 | faiss-cpu==1.7.2 4 | treelite==2.4.0 5 | treelite_runtime==2.4.0 6 | torch~=1.12 7 | lightgbm==3.3.2 8 | ipykernel 9 | -------------------------------------------------------------------------------- /requirements/test-gpu.txt: -------------------------------------------------------------------------------- 1 | -r test.txt 2 | 3 | faiss-gpu==1.7.2 4 | tritonclient==2.30.0 -------------------------------------------------------------------------------- /requirements/test.txt: -------------------------------------------------------------------------------- 1 | -r base.txt 2 | -r dev.txt 3 | 4 | # This contains common libraries for testing. 5 | 6 | # NOTE: You should pip install requirements/test-[cpu|gpu].txt for device-specific test 7 | # requirements, which will include the dependencies defined in this file. 8 | 9 | pytest>=5 10 | pytest-cov>=2 11 | scikit-learn<1.2 12 | asvdb@git+https://github.com/rapidsai/asvdb.git 13 | testbook==0.4.2 14 | 15 | # packages necessary to run tests and push PRs 16 | tritonclient 17 | feast==0.31 18 | xgboost==1.6.2 19 | implicit==0.6.0 20 | 21 | merlin-models[tensorflow,pytorch,transformers]@git+https://github.com/NVIDIA-Merlin/models.git 22 | 23 | # TODO: do we need more of these? 24 | # https://github.com/NVIDIA-Merlin/Merlin/blob/a1cc48fe23c4dfc627423168436f26ef7e028204/ci/dockerfile.ci#L13-L18 25 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length = 100 3 | exclude = build,.eggs,*_pb2.py 4 | ignore = E203,W503 5 | per-file-ignores = 6 | examples/criteo_benchmark.py:E402 7 | examples/dataloader_bench.py:E402 8 | 9 | [flake8_nb] 10 | max-line-length = 120 11 | ignore = E203,E402,W503 12 | 13 | [pydocstyle] 14 | ignore = D100,D102,D103,D104,D105,D107,D203,D205,D211,D212,D213,D400,D401,D413,D415 15 | 16 | [codespell] 17 | skip = *pb2.py,./.git,./.github,./bench,./dist,./docs/build,.*egg-info.*,versioneer.py,*.csv,*.parquet 18 | ignore-words = ./ci/ignore_codespell_words.txt 19 | count = 20 | quiet-level = 3 21 | 22 | # See the docstring in versioneer.py for instructions. Note that you must 23 | # re-run 'versioneer.py setup' after changing this section, and commit the 24 | # resulting files. 25 | 26 | [versioneer] 27 | VCS = git 28 | style = pep440 29 | versionfile_source = merlin/systems/_version.py 30 | versionfile_build = merlin/systems/_version.py 31 | tag_prefix = v 32 | parentdir_prefix = merlin-systems- 33 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import codecs 17 | import os 18 | import sys 19 | 20 | from setuptools import find_namespace_packages, setup 21 | 22 | try: 23 | import versioneer 24 | except ImportError: 25 | # we have a versioneer.py file living in the same directory as this file, but 26 | # if we're using pep 517/518 to build from pyproject.toml its not going to find it 27 | # https://github.com/python-versioneer/python-versioneer/issues/193#issue-408237852 28 | # make this work by adding this directory to the python path 29 | sys.path.append(os.path.dirname(os.path.realpath(__file__))) 30 | import versioneer 31 | 32 | 33 | def read_requirements(req_path, filename): 34 | base = os.path.abspath(os.path.dirname(__file__)) 35 | with codecs.open(os.path.join(base, req_path, filename), "rb", "utf-8") as f: 36 | lineiter = (line.strip() for line in f) 37 | packages = [] 38 | for line in lineiter: 39 | if line: 40 | if line.startswith("-r"): 41 | filename = line.replace("-r", "").strip() 42 | packages.extend(read_requirements(req_path, filename)) 43 | elif not line.startswith("#"): 44 | packages.append(line) 45 | return packages 46 | 47 | 48 | requirements = { 49 | "cpu": read_requirements("requirements", "base.txt"), 50 | "gpu": read_requirements("requirements", "gpu.txt"), 51 | } 52 | dev_requirements = { 53 | "dev": read_requirements("requirements", "dev.txt"), 54 | "test": read_requirements("requirements", "test.txt"), 55 | "test-cpu": read_requirements("requirements", "test-cpu.txt"), 56 | "test-gpu": read_requirements("requirements", "test-gpu.txt"), 57 | "docs": read_requirements("requirements", "docs.txt"), 58 | } 59 | 60 | with open("README.md", encoding="utf8") as readme_file: 61 | long_description = readme_file.read() 62 | 63 | setup( 64 | name="merlin-systems", 65 | version=versioneer.get_version(), 66 | packages=find_namespace_packages(include=["merlin*"]), 67 | url="https://github.com/NVIDIA-Merlin/systems", 68 | author="NVIDIA Corporation", 69 | license="Apache 2.0", 70 | long_description=long_description, 71 | long_description_content_type="text/markdown", 72 | classifiers=[ 73 | "Development Status :: 4 - Beta", 74 | "Programming Language :: Python :: 3", 75 | "Programming Language :: Python :: 3.8", 76 | "Programming Language :: Python :: 3.9", 77 | "Programming Language :: Python :: 3.10", 78 | "Programming Language :: Python :: 3 :: Only", 79 | "Intended Audience :: Developers", 80 | "License :: OSI Approved :: Apache Software License", 81 | "Topic :: Software Development :: Libraries", 82 | "Topic :: Scientific/Engineering", 83 | ], 84 | zip_safe=False, 85 | python_requires=">=3.8", 86 | install_requires=requirements["cpu"], 87 | extras_require={ 88 | **requirements, 89 | **dev_requirements, 90 | }, 91 | cmdclass=versioneer.get_cmdclass(), 92 | ) 93 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/integration/t4r/test_pytorch_backend.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import shutil 18 | 19 | import pytest 20 | 21 | np = pytest.importorskip("numpy") 22 | torch = pytest.importorskip("torch") 23 | t4r = pytest.importorskip("transformers4rec") 24 | tr = pytest.importorskip("transformers4rec.torch") 25 | 26 | triton = pytest.importorskip("merlin.systems.triton") 27 | data_conversions = pytest.importorskip("merlin.systems.triton.conversions") 28 | 29 | tritonclient = pytest.importorskip("tritonclient") 30 | grpcclient = pytest.importorskip("tritonclient.grpc") 31 | 32 | from merlin.core.compat import HAS_GPU # noqa 33 | from merlin.core.dispatch import make_df # noqa 34 | from merlin.systems.dag import Ensemble # noqa 35 | from merlin.systems.dag.ops.pytorch import PredictPyTorch # noqa 36 | from merlin.systems.triton.conversions import match_representations # noqa 37 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver # noqa 38 | 39 | TRITON_SERVER_PATH = shutil.which("tritonserver") 40 | 41 | 42 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 43 | @pytest.mark.skipif(not HAS_GPU, reason="GPU Device required for test") 44 | def test_serve_t4r_with_torchscript(tmpdir): 45 | # =========================================== 46 | # Generate training data 47 | # =========================================== 48 | 49 | min_session_len = 5 50 | max_session_len = 20 51 | torch_yoochoose_like = tr.data.tabular_sequence_testing_data.torch_synthetic_data( 52 | num_rows=100, 53 | min_session_length=min_session_len, 54 | max_session_length=max_session_len, 55 | device="cuda", 56 | ) 57 | t4r_yoochoose_schema = t4r.data.tabular_sequence_testing_data.schema 58 | 59 | # =========================================== 60 | # Build, train, test, and JIT the model 61 | # =========================================== 62 | 63 | input_module = t4r.torch.TabularSequenceFeatures.from_schema( 64 | t4r_yoochoose_schema, 65 | max_sequence_length=20, 66 | d_output=64, 67 | masking="causal", 68 | ) 69 | prediction_task = t4r.torch.NextItemPredictionTask(weight_tying=True) 70 | transformer_config = t4r.config.transformer.XLNetConfig.build( 71 | d_model=64, n_head=8, n_layer=2, total_seq_length=20 72 | ) 73 | model = transformer_config.to_torch_model(input_module, prediction_task) 74 | model = model.cuda() 75 | 76 | _ = model(torch_yoochoose_like) 77 | 78 | model.eval() 79 | 80 | example_inputs = match_representations(model.input_schema, torch_yoochoose_like) 81 | traced_model = torch.jit.trace(model, example_inputs, strict=True) 82 | assert isinstance(traced_model, torch.jit.TopLevelTracedModule) 83 | assert torch.allclose( 84 | model(example_inputs), 85 | traced_model(example_inputs), 86 | ) 87 | 88 | # =========================================== 89 | # Build a simple Ensemble graph 90 | # =========================================== 91 | 92 | input_schema = model.input_schema 93 | output_schema = model.output_schema 94 | 95 | torch_op = input_schema.column_names >> PredictPyTorch( 96 | traced_model, input_schema, output_schema 97 | ) 98 | 99 | ensemble = Ensemble(torch_op, input_schema) 100 | ens_config, node_configs = ensemble.export(str(tmpdir)) 101 | 102 | # =========================================== 103 | # Create Request Data 104 | # =========================================== 105 | 106 | request_data = tr.data.tabular_sequence_testing_data.torch_synthetic_data( 107 | num_rows=40, 108 | min_session_length=4, 109 | max_session_length=10, 110 | device="cuda", 111 | ) 112 | 113 | df_cols = {} 114 | for name, tensor in request_data.items(): 115 | if name in input_schema.column_names: 116 | dtype = input_schema[name].dtype 117 | 118 | df_cols[name] = tensor.cpu().numpy().astype(dtype.name) 119 | if len(tensor.shape) > 1: 120 | df_cols[name] = list(df_cols[name]) 121 | 122 | df = make_df(df_cols) 123 | 124 | # =========================================== 125 | # Send request to Triton and check response 126 | # =========================================== 127 | triton_response = run_ensemble_on_tritonserver( 128 | tmpdir, input_schema, df, output_schema.column_names, "executor_model" 129 | ) 130 | 131 | assert triton_response 132 | 133 | preds_triton = triton_response[output_schema.column_names[0]] 134 | 135 | preds_model = model(request_data).cpu().detach().numpy() 136 | 137 | np.testing.assert_allclose(preds_triton, preds_model) 138 | -------------------------------------------------------------------------------- /tests/integration/tf/test_transformer_model.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import shutil 18 | 19 | import pytest 20 | 21 | tf = pytest.importorskip("tensorflow") 22 | 23 | triton = pytest.importorskip("merlin.systems.triton") 24 | 25 | tritonclient = pytest.importorskip("tritonclient") 26 | grpcclient = pytest.importorskip("tritonclient.grpc") 27 | 28 | import merlin.models.tf as mm # noqa 29 | from merlin.datasets.synthetic import generate_data # noqa 30 | from merlin.io import Dataset # noqa 31 | from merlin.schema import Tags # noqa 32 | from merlin.systems.dag import Ensemble # noqa 33 | from merlin.systems.dag.ops.tensorflow import PredictTensorflow # noqa 34 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver # noqa 35 | 36 | TRITON_SERVER_PATH = shutil.which("tritonserver") 37 | 38 | 39 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 40 | def test_serve_tf_session_based_with_libtensorflow(tmpdir): 41 | 42 | # =========================================== 43 | # Generate training data 44 | # =========================================== 45 | 46 | train = generate_data("sequence-testing", num_rows=100) 47 | 48 | # =========================================== 49 | # Build and train the model 50 | # =========================================== 51 | 52 | seq_schema = train.schema.select_by_tag(Tags.SEQUENCE).select_by_tag(Tags.CATEGORICAL) 53 | 54 | target = train.schema.select_by_tag(Tags.ITEM_ID).column_names[0] 55 | predict_last = mm.SequencePredictLast(schema=seq_schema, target=target) 56 | 57 | input_schema = seq_schema 58 | output_schema = seq_schema.select_by_name(target) 59 | 60 | train = Dataset(train.to_ddf(columns=input_schema.column_names).compute()) 61 | train.schema = input_schema 62 | loader = mm.Loader(train, batch_size=16, shuffle=False) 63 | 64 | d_model = 48 65 | query_encoder = mm.Encoder( 66 | mm.InputBlockV2( 67 | input_schema, 68 | embeddings=mm.Embeddings( 69 | input_schema.select_by_tag(Tags.CATEGORICAL), sequence_combiner=None 70 | ), 71 | ), 72 | mm.MLPBlock([d_model]), 73 | mm.GPT2Block(d_model=d_model, n_head=2, n_layer=2), 74 | tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=1)), 75 | ) 76 | 77 | model = mm.RetrievalModelV2( 78 | query=query_encoder, 79 | output=mm.ContrastiveOutput(output_schema, negative_samplers="in-batch"), 80 | ) 81 | 82 | model.compile(metrics={}) 83 | model.fit(loader, epochs=1, pre=predict_last) 84 | 85 | # =========================================== 86 | # Build a simple Ensemble graph 87 | # =========================================== 88 | tf_op = input_schema.column_names >> PredictTensorflow(model.query_encoder) 89 | 90 | ensemble = Ensemble(tf_op, input_schema) 91 | ens_config, node_configs = ensemble.export(str(tmpdir)) 92 | 93 | # =========================================== 94 | # Create Request Data 95 | # =========================================== 96 | 97 | data = generate_data("sequence-testing", num_rows=1) 98 | request_df = data.compute() 99 | 100 | # =========================================== 101 | # Send request to Triton and check response 102 | # =========================================== 103 | response = run_ensemble_on_tritonserver( 104 | tmpdir, input_schema, request_df, ["output_1"], node_configs[0].name 105 | ) 106 | 107 | assert response 108 | assert len(response["output_1"][0]) == d_model 109 | -------------------------------------------------------------------------------- /tests/test_passing.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2023, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import pytest 17 | 18 | 19 | # This test exists to prevent pre- and post-merge test runs from finding no tests 20 | # and then returning a non-zero exit code. It always passes by design. 21 | @pytest.mark.premerge 22 | @pytest.mark.postmerge 23 | def test_passing(): 24 | pass 25 | -------------------------------------------------------------------------------- /tests/unit/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/ops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-Merlin/systems/a19d3113e296c69699267bc23f991e5fc4db4ef6/tests/unit/systems/dag/ops/__init__.py -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/fil/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/fil/test_lightgbm.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import pytest 4 | 5 | from merlin.dag import ColumnSelector 6 | from merlin.dag.runtime import Runtime 7 | from merlin.dtypes.shape import Shape 8 | from merlin.schema import ColumnSchema, Schema 9 | from merlin.systems.dag.ensemble import Ensemble 10 | from merlin.systems.dag.ops.fil import PredictForest 11 | from merlin.table import TensorTable 12 | 13 | sklearn_datasets = pytest.importorskip("sklearn.datasets") 14 | lightgbm = pytest.importorskip("lightgbm") 15 | export = pytest.importorskip("merlin.systems.dag.ensemble") 16 | 17 | 18 | @pytest.mark.parametrize("runtime", [Runtime()]) 19 | def test_lightgbm_regressor_forest_inference(runtime, tmpdir): 20 | rows = 200 21 | num_features = 16 22 | X, y = sklearn_datasets.make_regression( 23 | n_samples=rows, 24 | n_features=num_features, 25 | n_informative=num_features // 3, 26 | random_state=0, 27 | ) 28 | feature_names = [str(i) for i in range(num_features)] 29 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 30 | for column in df.columns: 31 | df[column] = np.log(df[column] + 1).fillna(0.5) 32 | 33 | # Fit GBDT Model 34 | model = lightgbm.LGBMRegressor() 35 | model.fit(X, y) 36 | 37 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 38 | input_schema = Schema(input_column_schemas) 39 | selector = ColumnSelector(feature_names) 40 | 41 | triton_chain = selector >> PredictForest(model, input_schema) 42 | 43 | ensemble = Ensemble(triton_chain, input_schema) 44 | 45 | request_df = df[:5] 46 | 47 | tensor_table = TensorTable.from_df(request_df) 48 | response = ensemble.transform(tensor_table, runtime=runtime) 49 | 50 | assert response["output__0"].shape == Shape((5,)) 51 | 52 | 53 | @pytest.mark.parametrize("runtime", [Runtime()]) 54 | def test_lightgbm_classify_forest_inference(runtime, tmpdir): 55 | rows = 200 56 | num_features = 16 57 | X, y = sklearn_datasets.make_classification( 58 | n_samples=rows, 59 | n_features=num_features, 60 | n_informative=num_features // 3, 61 | random_state=0, 62 | ) 63 | feature_names = [str(i) for i in range(num_features)] 64 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 65 | for column in df.columns: 66 | df[column] = np.log(df[column] + 1).fillna(0.5) 67 | 68 | # Fit GBDT Model 69 | model = lightgbm.LGBMClassifier() 70 | model.fit(X, y) 71 | 72 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 73 | input_schema = Schema(input_column_schemas) 74 | selector = ColumnSelector(feature_names) 75 | 76 | triton_chain = selector >> PredictForest(model, input_schema) 77 | 78 | ensemble = Ensemble(triton_chain, input_schema) 79 | 80 | request_df = df[:5] 81 | 82 | tensor_table = TensorTable.from_df(request_df) 83 | response = ensemble.transform(tensor_table, runtime=runtime) 84 | 85 | assert response["output__0"].shape == Shape((5,)) 86 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/fil/test_sklearn.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import pytest 4 | 5 | from merlin.dag import ColumnSelector 6 | from merlin.dag.runtime import Runtime 7 | from merlin.dtypes.shape import Shape 8 | from merlin.schema import ColumnSchema, Schema 9 | from merlin.systems.dag.ensemble import Ensemble 10 | from merlin.systems.dag.ops.fil import PredictForest 11 | from merlin.table import TensorTable 12 | 13 | sklearn_datasets = pytest.importorskip("sklearn.datasets") 14 | sklearn_ensemble = pytest.importorskip("sklearn.ensemble") 15 | export = pytest.importorskip("merlin.systems.dag.ensemble") 16 | 17 | 18 | @pytest.mark.parametrize("runtime", [Runtime()]) 19 | def test_sklearn_regressor_forest_inference(runtime, tmpdir): 20 | rows = 200 21 | num_features = 16 22 | X, y = sklearn_datasets.make_regression( 23 | n_samples=rows, 24 | n_features=num_features, 25 | n_informative=num_features // 3, 26 | random_state=0, 27 | ) 28 | feature_names = [str(i) for i in range(num_features)] 29 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 30 | for column in df.columns: 31 | df[column] = np.log(df[column] + 1).fillna(0.5) 32 | 33 | # Fit GBDT Model 34 | model = sklearn_ensemble.RandomForestRegressor() 35 | model.fit(X, y) 36 | 37 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 38 | input_schema = Schema(input_column_schemas) 39 | selector = ColumnSelector(feature_names) 40 | 41 | triton_chain = selector >> PredictForest(model, input_schema) 42 | 43 | ensemble = Ensemble(triton_chain, input_schema) 44 | 45 | request_df = df[:5] 46 | 47 | response = None 48 | 49 | tensor_table = TensorTable.from_df(request_df) 50 | response = ensemble.transform(tensor_table, runtime=runtime) 51 | 52 | assert response["output__0"].shape == Shape((5,)) 53 | 54 | 55 | @pytest.mark.parametrize("runtime", [Runtime()]) 56 | def test_sklearn_classify_forest_inference(runtime, tmpdir): 57 | rows = 200 58 | num_features = 16 59 | X, y = sklearn_datasets.make_classification( 60 | n_samples=rows, 61 | n_features=num_features, 62 | n_informative=num_features // 3, 63 | random_state=0, 64 | ) 65 | feature_names = [str(i) for i in range(num_features)] 66 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 67 | for column in df.columns: 68 | df[column] = np.log(df[column] + 1).fillna(0.5) 69 | 70 | # Fit GBDT Model 71 | model = sklearn_ensemble.RandomForestClassifier() 72 | model.fit(X, y) 73 | 74 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 75 | input_schema = Schema(input_column_schemas) 76 | selector = ColumnSelector(feature_names) 77 | 78 | triton_chain = selector >> PredictForest(model, input_schema) 79 | 80 | ensemble = Ensemble(triton_chain, input_schema) 81 | 82 | request_df = df[:5] 83 | 84 | tensor_table = TensorTable.from_df(request_df) 85 | response = ensemble.transform(tensor_table, runtime=runtime) 86 | 87 | assert response["output__0"].shape == Shape((5,)) 88 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/fil/test_xgboost.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import pytest 4 | 5 | from merlin.core.compat import HAS_GPU 6 | from merlin.dag import ColumnSelector 7 | from merlin.dag.runtime import Runtime 8 | from merlin.dtypes.shape import Shape 9 | from merlin.schema import ColumnSchema, Schema 10 | from merlin.systems.dag.ensemble import Ensemble 11 | from merlin.systems.dag.ops.fil import PredictForest 12 | from merlin.table import TensorTable 13 | 14 | sklearn_datasets = pytest.importorskip("sklearn.datasets") 15 | xgboost = pytest.importorskip("xgboost") 16 | export = pytest.importorskip("merlin.systems.dag.ensemble") 17 | 18 | 19 | @pytest.mark.parametrize("runtime", [Runtime()]) 20 | @pytest.mark.skipif(not HAS_GPU, reason="no gpu detected") 21 | def test_xgboost_regressor_forest_inference(runtime, tmpdir): 22 | rows = 200 23 | num_features = 16 24 | X, y = sklearn_datasets.make_regression( 25 | n_samples=rows, 26 | n_features=num_features, 27 | n_informative=num_features // 3, 28 | random_state=0, 29 | ) 30 | feature_names = [str(i) for i in range(num_features)] 31 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 32 | for column in df.columns: 33 | df[column] = np.log(df[column] + 1).fillna(0.5) 34 | 35 | # Fit GBDT Model 36 | model = xgboost.XGBRegressor() 37 | model.fit(X, y) 38 | 39 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 40 | input_schema = Schema(input_column_schemas) 41 | selector = ColumnSelector(feature_names) 42 | 43 | triton_chain = selector >> PredictForest(model, input_schema) 44 | 45 | ensemble = Ensemble(triton_chain, input_schema) 46 | 47 | request_df = df[:5] 48 | 49 | tensor_table = TensorTable.from_df(request_df) 50 | response = ensemble.transform(tensor_table, runtime=runtime) 51 | 52 | assert response["output__0"].shape == Shape((5,)) 53 | 54 | 55 | @pytest.mark.skipif(not HAS_GPU, reason="no gpu detected") 56 | @pytest.mark.parametrize("runtime", [Runtime()]) 57 | def test_xgboost_classify_forest_inference(runtime, tmpdir): 58 | rows = 200 59 | num_features = 16 60 | X, y = sklearn_datasets.make_classification( 61 | n_samples=rows, 62 | n_features=num_features, 63 | n_informative=num_features // 3, 64 | random_state=0, 65 | ) 66 | feature_names = [str(i) for i in range(num_features)] 67 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 68 | for column in df.columns: 69 | df[column] = np.log(df[column] + 1).fillna(0.5) 70 | 71 | # Fit GBDT Model 72 | model = xgboost.XGBClassifier() 73 | model.fit(X, y) 74 | 75 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 76 | input_schema = Schema(input_column_schemas) 77 | selector = ColumnSelector(feature_names) 78 | 79 | triton_chain = selector >> PredictForest(model, input_schema) 80 | 81 | ensemble = Ensemble(triton_chain, input_schema) 82 | 83 | request_df = df[:5] 84 | 85 | tensor_table = TensorTable.from_df(request_df) 86 | response = ensemble.transform(tensor_table, runtime=runtime) 87 | 88 | assert response["output__0"].shape == Shape((5,)) 89 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/nvtabular/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/nvtabular/test_ensemble.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import numpy as np 17 | import pytest 18 | 19 | from merlin.dag.runtime import Runtime 20 | from merlin.table import TensorTable 21 | from nvtabular import Workflow 22 | from nvtabular import ops as wf_ops 23 | 24 | ensemble = pytest.importorskip("merlin.systems.dag.ensemble") 25 | workflow_op = pytest.importorskip("merlin.systems.dag.ops.workflow") 26 | 27 | 28 | @pytest.mark.parametrize("engine", ["parquet"]) 29 | @pytest.mark.parametrize("runtime", [Runtime()]) 30 | def test_workflow_op_serving_triton(tmpdir, dataset, engine, runtime): 31 | input_columns = ["x", "y", "id"] 32 | 33 | # NVT 34 | workflow_ops = input_columns >> wf_ops.Rename(postfix="_nvt") 35 | workflow = Workflow(workflow_ops) 36 | workflow.fit(dataset) 37 | 38 | # Triton 39 | triton_op = "*" >> workflow_op.TransformWorkflow( 40 | workflow, 41 | conts=["x_nvt", "y_nvt"], 42 | cats=["id_nvt"], 43 | ) 44 | 45 | wkflow_ensemble = ensemble.Ensemble(triton_op, workflow.input_schema) 46 | 47 | input_data = {} 48 | for col_name, col_schema in workflow.input_schema.column_schemas.items(): 49 | col_dtype = col_schema.dtype 50 | input_data[col_name] = np.array([2, 3, 4]).astype(col_dtype.to_numpy) 51 | table = TensorTable(input_data) 52 | response = wkflow_ensemble.transform(table, runtime=runtime) 53 | 54 | for col_name in workflow.output_schema.column_names: 55 | assert response[col_name].shape.dims[0] == table[col_name.split("_")[0]].shape.dims[0] 56 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/tensorflow/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/torch/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | pytest.importorskip("torch") 20 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local/ops/torch/test_op.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | from typing import Dict 17 | 18 | import numpy as np 19 | import pytest 20 | 21 | from merlin.dag.runtime import Runtime 22 | from merlin.schema import ColumnSchema, Schema 23 | from merlin.systems.dag.ensemble import Ensemble 24 | from merlin.table import TensorTable 25 | 26 | torch = pytest.importorskip("torch") 27 | ptorch_op = pytest.importorskip("merlin.systems.dag.ops.pytorch") 28 | 29 | 30 | class CustomModel(torch.nn.Module): 31 | def __init__(self): 32 | super().__init__() 33 | self.linear = torch.nn.Linear(3, 1) 34 | 35 | def forward(self, input_dict: Dict[str, torch.Tensor]): 36 | linear_out = self.linear(input_dict["input"].to(self.linear.weight.device)) 37 | 38 | return linear_out 39 | 40 | 41 | model = CustomModel() 42 | model_scripted = torch.jit.script(model) 43 | 44 | model_input_schema = Schema( 45 | [ 46 | ColumnSchema( 47 | "input", 48 | properties={"value_count": {"min": 3, "max": 3}}, 49 | dtype=np.float32, 50 | is_list=True, 51 | is_ragged=False, 52 | ) 53 | ] 54 | ) 55 | model_output_schema = Schema([ColumnSchema("OUTPUT__0", dtype=np.float32)]) 56 | 57 | 58 | @pytest.mark.parametrize("torchscript", [True]) 59 | @pytest.mark.parametrize("use_path", [True, False]) 60 | @pytest.mark.parametrize("runtime", [(Runtime())]) 61 | def test_pytorch_op_serving(tmpdir, use_path, torchscript, runtime): 62 | model_path = str(tmpdir / "model.pt") 63 | 64 | model_to_use = model_scripted if torchscript else model 65 | model_or_path = model_path if use_path else model_to_use 66 | 67 | if use_path: 68 | try: 69 | # jit-compiled version of a model 70 | model_to_use.save(model_path) 71 | except AttributeError: 72 | # non-jit-compiled version of a model 73 | torch.save(model_to_use, model_path) 74 | 75 | predictions = ["input"] >> ptorch_op.PredictPyTorch( 76 | model_or_path, model_input_schema, model_output_schema 77 | ) 78 | ensemble = Ensemble(predictions, model_input_schema) 79 | 80 | input_data = {"input": np.array([[2.0, 3.0, 4.0], [4.0, 8.0, 1.0]]).astype(np.float32)} 81 | 82 | inputs = TensorTable(input_data) 83 | response = ensemble.transform(inputs, runtime=runtime) 84 | 85 | assert response["OUTPUT__0"].values.shape[0] == input_data["input"].shape[0] 86 | 87 | 88 | @pytest.mark.parametrize("torchscript", [True]) 89 | @pytest.mark.parametrize("use_path", [True, False]) 90 | @pytest.mark.parametrize("runtime", [(Runtime())]) 91 | def test_pytorch_op_serving_python(tmpdir, use_path, torchscript, runtime): 92 | model_path = str(tmpdir / "model.pt") 93 | 94 | model_to_use = model_scripted if torchscript else model 95 | model_or_path = model_path if use_path else model_to_use 96 | 97 | if use_path: 98 | try: 99 | # jit-compiled version of a model 100 | model_to_use.save(model_path) 101 | except AttributeError: 102 | # non-jit-compiled version of a model 103 | torch.save(model_to_use, model_path) 104 | 105 | predictions = ["input"] >> ptorch_op.PredictPyTorch( 106 | model_or_path, model_input_schema, model_output_schema 107 | ) 108 | ensemble = Ensemble(predictions, model_input_schema) 109 | 110 | input_data = {"input": np.array([[2.0, 3.0, 4.0], [4.0, 8.0, 1.0]]).astype(np.float32)} 111 | 112 | inputs = TensorTable(input_data) 113 | response = ensemble.transform(inputs, runtime=runtime) 114 | assert response["OUTPUT__0"].values.shape[0] == input_data["input"].shape[0] 115 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local_runtime/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local_runtime/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/local_runtime/ops/torch/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/fil/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/fil/test_lightgbm_triton.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | import numpy as np 4 | import pandas as pd 5 | import pytest 6 | 7 | from merlin.dag import ColumnSelector 8 | from merlin.schema import ColumnSchema, Schema 9 | from merlin.systems.dag.ensemble import Ensemble 10 | from merlin.systems.dag.ops.fil import PredictForest 11 | from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime 12 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver 13 | 14 | sklearn_datasets = pytest.importorskip("sklearn.datasets") 15 | lightgbm = pytest.importorskip("lightgbm") 16 | triton = pytest.importorskip("merlin.systems.triton") 17 | export = pytest.importorskip("merlin.systems.dag.ensemble") 18 | 19 | TRITON_SERVER_PATH = shutil.which("tritonserver") 20 | 21 | 22 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 23 | @pytest.mark.parametrize( 24 | ["runtime", "model_name", "expected_model_name"], 25 | [ 26 | (TritonExecutorRuntime(), None, "executor_model"), 27 | ], 28 | ) 29 | def test_lightgbm_regressor_forest_inference(runtime, model_name, expected_model_name, tmpdir): 30 | rows = 200 31 | num_features = 16 32 | X, y = sklearn_datasets.make_regression( 33 | n_samples=rows, 34 | n_features=num_features, 35 | n_informative=num_features // 3, 36 | random_state=0, 37 | ) 38 | feature_names = [str(i) for i in range(num_features)] 39 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 40 | for column in df.columns: 41 | df[column] = np.log(df[column] + 1).fillna(0.5) 42 | 43 | # Fit GBDT Model 44 | model = lightgbm.LGBMRegressor() 45 | model.fit(X, y) 46 | 47 | request_df = df[:5] 48 | 49 | input_column_schemas = [ 50 | ColumnSchema(col, dtype=np.float32, dims=request_df[col].shape) for col in feature_names 51 | ] 52 | input_schema = Schema(input_column_schemas) 53 | selector = ColumnSelector(feature_names) 54 | 55 | triton_chain = selector >> PredictForest(model, input_schema) 56 | 57 | ensemble = Ensemble(triton_chain, input_schema) 58 | 59 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime, name=model_name) 60 | 61 | assert ensemble_config.name == expected_model_name 62 | 63 | response = run_ensemble_on_tritonserver( 64 | str(tmpdir), input_schema, request_df, ["output__0"], ensemble_config.name 65 | ) 66 | assert response["output__0"].shape == (5,) 67 | 68 | 69 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 70 | @pytest.mark.parametrize( 71 | ["runtime", "model_name", "expected_model_name"], 72 | [ 73 | (TritonExecutorRuntime(), None, "executor_model"), 74 | ], 75 | ) 76 | def test_lightgbm_classify_forest_inference(runtime, model_name, expected_model_name, tmpdir): 77 | rows = 200 78 | num_features = 16 79 | X, y = sklearn_datasets.make_classification( 80 | n_samples=rows, 81 | n_features=num_features, 82 | n_informative=num_features // 3, 83 | random_state=0, 84 | ) 85 | feature_names = [str(i) for i in range(num_features)] 86 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 87 | for column in df.columns: 88 | df[column] = np.log(df[column] + 1).fillna(0.5) 89 | 90 | # Fit GBDT Model 91 | model = lightgbm.LGBMClassifier() 92 | model.fit(X, y) 93 | 94 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 95 | input_schema = Schema(input_column_schemas) 96 | selector = ColumnSelector(feature_names) 97 | 98 | triton_chain = selector >> PredictForest(model, input_schema) 99 | 100 | ensemble = Ensemble(triton_chain, input_schema) 101 | 102 | request_df = df[:5] 103 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime, name=model_name) 104 | 105 | assert ensemble_config.name == expected_model_name 106 | 107 | response = run_ensemble_on_tritonserver( 108 | str(tmpdir), input_schema, request_df, ["output__0"], ensemble_config.name 109 | ) 110 | assert response["output__0"].shape == (5,) 111 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/fil/test_sklearn_triton.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | import numpy as np 4 | import pandas as pd 5 | import pytest 6 | 7 | from merlin.dag import ColumnSelector 8 | from merlin.schema import ColumnSchema, Schema 9 | from merlin.systems.dag.ensemble import Ensemble 10 | from merlin.systems.dag.ops.fil import PredictForest 11 | from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime 12 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver 13 | 14 | sklearn_datasets = pytest.importorskip("sklearn.datasets") 15 | sklearn_ensemble = pytest.importorskip("sklearn.ensemble") 16 | triton = pytest.importorskip("merlin.systems.triton") 17 | export = pytest.importorskip("merlin.systems.dag.ensemble") 18 | 19 | TRITON_SERVER_PATH = shutil.which("tritonserver") 20 | 21 | 22 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 23 | @pytest.mark.parametrize( 24 | ["runtime", "model_name", "expected_model_name"], 25 | [ 26 | (TritonExecutorRuntime(), None, "executor_model"), 27 | ], 28 | ) 29 | def test_sklearn_regressor_forest_inference(runtime, model_name, expected_model_name, tmpdir): 30 | rows = 200 31 | num_features = 16 32 | X, y = sklearn_datasets.make_regression( 33 | n_samples=rows, 34 | n_features=num_features, 35 | n_informative=num_features // 3, 36 | random_state=0, 37 | ) 38 | feature_names = [str(i) for i in range(num_features)] 39 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 40 | for column in df.columns: 41 | df[column] = np.log(df[column] + 1).fillna(0.5) 42 | 43 | # Fit GBDT Model 44 | model = sklearn_ensemble.RandomForestRegressor() 45 | model.fit(X, y) 46 | 47 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 48 | input_schema = Schema(input_column_schemas) 49 | selector = ColumnSelector(feature_names) 50 | 51 | triton_chain = selector >> PredictForest(model, input_schema) 52 | 53 | ensemble = Ensemble(triton_chain, input_schema) 54 | 55 | request_df = df[:5] 56 | 57 | response = None 58 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime, name=model_name) 59 | assert ensemble_config.name == expected_model_name 60 | 61 | response = run_ensemble_on_tritonserver( 62 | str(tmpdir), input_schema, request_df, ["output__0"], ensemble_config.name 63 | ) 64 | assert response["output__0"].shape == (5,) 65 | 66 | 67 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 68 | @pytest.mark.parametrize( 69 | ["runtime", "model_name", "expected_model_name"], 70 | [ 71 | (TritonExecutorRuntime(), None, "executor_model"), 72 | ], 73 | ) 74 | def test_sklearn_classify_forest_inference(runtime, model_name, expected_model_name, tmpdir): 75 | rows = 200 76 | num_features = 16 77 | X, y = sklearn_datasets.make_classification( 78 | n_samples=rows, 79 | n_features=num_features, 80 | n_informative=num_features // 3, 81 | random_state=0, 82 | ) 83 | feature_names = [str(i) for i in range(num_features)] 84 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 85 | for column in df.columns: 86 | df[column] = np.log(df[column] + 1).fillna(0.5) 87 | 88 | # Fit GBDT Model 89 | model = sklearn_ensemble.RandomForestClassifier() 90 | model.fit(X, y) 91 | 92 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 93 | input_schema = Schema(input_column_schemas) 94 | selector = ColumnSelector(feature_names) 95 | 96 | triton_chain = selector >> PredictForest(model, input_schema) 97 | 98 | ensemble = Ensemble(triton_chain, input_schema) 99 | 100 | request_df = df[:5] 101 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime, name=model_name) 102 | 103 | assert ensemble_config.name == expected_model_name 104 | 105 | response = run_ensemble_on_tritonserver( 106 | str(tmpdir), input_schema, request_df, ["output__0"], ensemble_config.name 107 | ) 108 | assert response["output__0"].shape == (5,) 109 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/fil/test_xgboost_triton.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | import numpy as np 4 | import pandas as pd 5 | import pytest 6 | 7 | from merlin.dag import ColumnSelector 8 | from merlin.schema import ColumnSchema, Schema 9 | from merlin.systems.dag.ensemble import Ensemble 10 | from merlin.systems.dag.ops.fil import PredictForest 11 | from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime 12 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver 13 | 14 | sklearn_datasets = pytest.importorskip("sklearn.datasets") 15 | xgboost = pytest.importorskip("xgboost") 16 | triton = pytest.importorskip("merlin.systems.triton") 17 | export = pytest.importorskip("merlin.systems.dag.ensemble") 18 | 19 | TRITON_SERVER_PATH = shutil.which("tritonserver") 20 | 21 | 22 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 23 | @pytest.mark.parametrize( 24 | ["runtime", "model_name", "expected_model_name"], 25 | [ 26 | (TritonExecutorRuntime(), None, "executor_model"), 27 | ], 28 | ) 29 | def test_xgboost_regressor_forest_inference(runtime, model_name, expected_model_name, tmpdir): 30 | rows = 200 31 | num_features = 16 32 | X, y = sklearn_datasets.make_regression( 33 | n_samples=rows, 34 | n_features=num_features, 35 | n_informative=num_features // 3, 36 | random_state=0, 37 | ) 38 | feature_names = [str(i) for i in range(num_features)] 39 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 40 | for column in df.columns: 41 | df[column] = np.log(df[column] + 1).fillna(0.5) 42 | 43 | # Fit GBDT Model 44 | model = xgboost.XGBRegressor() 45 | model.fit(X, y) 46 | 47 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 48 | input_schema = Schema(input_column_schemas) 49 | selector = ColumnSelector(feature_names) 50 | 51 | triton_chain = selector >> PredictForest(model, input_schema) 52 | 53 | ensemble = Ensemble(triton_chain, input_schema) 54 | 55 | request_df = df[:5] 56 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime, name=model_name) 57 | 58 | assert ensemble_config.name == expected_model_name 59 | 60 | response = run_ensemble_on_tritonserver( 61 | str(tmpdir), input_schema, request_df, ["output__0"], ensemble_config.name 62 | ) 63 | assert response["output__0"].shape == (5,) 64 | 65 | 66 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 67 | @pytest.mark.parametrize( 68 | ["runtime", "model_name", "expected_model_name"], 69 | [ 70 | (TritonExecutorRuntime(), None, "executor_model"), 71 | ], 72 | ) 73 | def test_xgboost_classify_forest_inference(runtime, model_name, expected_model_name, tmpdir): 74 | rows = 200 75 | num_features = 16 76 | X, y = sklearn_datasets.make_classification( 77 | n_samples=rows, 78 | n_features=num_features, 79 | n_informative=num_features // 3, 80 | random_state=0, 81 | ) 82 | feature_names = [str(i) for i in range(num_features)] 83 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 84 | for column in df.columns: 85 | df[column] = np.log(df[column] + 1).fillna(0.5) 86 | 87 | # Fit GBDT Model 88 | model = xgboost.XGBClassifier() 89 | model.fit(X, y) 90 | 91 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 92 | input_schema = Schema(input_column_schemas) 93 | selector = ColumnSelector(feature_names) 94 | 95 | triton_chain = selector >> PredictForest(model, input_schema) 96 | 97 | ensemble = Ensemble(triton_chain, input_schema) 98 | 99 | request_df = df[:5] 100 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime, name=model_name) 101 | 102 | assert ensemble_config.name == expected_model_name 103 | 104 | response = run_ensemble_on_tritonserver( 105 | str(tmpdir), input_schema, request_df, ["output__0"], ensemble_config.name 106 | ) 107 | assert response["output__0"].shape == (5,) 108 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/torch/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import pytest 17 | 18 | pytest.importorskip("torch") 19 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/workflow/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/ops/workflow/test_op.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import os 17 | import pathlib 18 | 19 | import pytest 20 | 21 | # this needs to be before any modules that import protobuf 22 | os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" 23 | 24 | from google.protobuf import text_format # noqa 25 | from tritonclient.grpc import model_config_pb2 as model_config # noqa 26 | 27 | from merlin.schema import Schema # noqa 28 | from nvtabular import Workflow # noqa 29 | from nvtabular import ops as wf_ops # noqa 30 | 31 | ensemble = pytest.importorskip("merlin.systems.dag.ensemble") 32 | wf_op = pytest.importorskip("merlin.systems.dag.ops.workflow") 33 | wf_triton_op = pytest.importorskip("merlin.systems.dag.runtimes.triton.ops.workflow") 34 | 35 | 36 | @pytest.mark.parametrize("engine", ["parquet"]) 37 | def test_workflow_op_validates_schemas(dataset, engine): 38 | input_columns = ["x", "y", "id"] 39 | request_schema = Schema(input_columns) 40 | 41 | # NVT 42 | workflow_ops = input_columns >> wf_ops.Rename(postfix="_nvt") 43 | workflow = Workflow(workflow_ops) 44 | workflow.fit(dataset) 45 | 46 | # Triton 47 | triton_ops = ["a", "b", "c"] >> wf_triton_op.TransformWorkflowTriton( 48 | wf_op.TransformWorkflow(workflow) 49 | ) 50 | 51 | with pytest.raises(ValueError) as exc_info: 52 | ensemble.Ensemble(triton_ops, request_schema) 53 | assert "Missing column" in str(exc_info.value) 54 | 55 | 56 | @pytest.mark.parametrize("engine", ["parquet"]) 57 | def test_workflow_op_exports_own_config(tmpdir, dataset, engine): 58 | input_columns = ["x", "y", "id"] 59 | 60 | # NVT 61 | workflow_ops = input_columns >> wf_ops.Rename(postfix="_nvt") 62 | workflow = Workflow(workflow_ops) 63 | workflow.fit(dataset) 64 | 65 | # Triton 66 | triton_op = wf_triton_op.TransformWorkflowTriton(wf_op.TransformWorkflow(workflow)) 67 | triton_op.export(tmpdir, None, None) 68 | 69 | # Export creates directory 70 | export_path = pathlib.Path(tmpdir) / triton_op.export_name 71 | assert export_path.exists() 72 | 73 | # Export creates the config file 74 | config_path = export_path / "config.pbtxt" 75 | assert config_path.exists() 76 | 77 | # Read the config file back in from proto 78 | with open(config_path, "rb") as f: 79 | config = model_config.ModelConfig() 80 | raw_config = f.read() 81 | parsed = text_format.Parse(raw_config, config) 82 | 83 | # The config file contents are correct 84 | assert parsed.name == triton_op.export_name 85 | assert parsed.backend == "python" 86 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/runtimes/triton/test_triton.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import random 18 | import shutil 19 | 20 | import numpy as np 21 | import pytest 22 | 23 | from merlin.schema import ColumnSchema, Schema 24 | from merlin.systems.dag.ensemble import Ensemble 25 | from merlin.systems.dag.ops.session_filter import FilterCandidates 26 | from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime 27 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver 28 | from merlin.table import TensorTable 29 | 30 | triton = pytest.importorskip("merlin.systems.triton") 31 | export = pytest.importorskip("merlin.systems.dag.ensemble") 32 | 33 | TRITON_SERVER_PATH = shutil.which("tritonserver") 34 | 35 | 36 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 37 | @pytest.mark.parametrize( 38 | ["runtime", "model_name", "expected_model_name"], 39 | [ 40 | (TritonExecutorRuntime(), None, "executor_model"), 41 | (TritonExecutorRuntime(), "triton_model", "triton_model"), 42 | ], 43 | ) 44 | def test_triton_runtime_export_and_run(runtime, model_name, expected_model_name, tmpdir): 45 | request_schema = Schema( 46 | [ 47 | ColumnSchema("candidate_ids", dtype=np.int32, dims=(None, 100)), 48 | ColumnSchema("movie_ids", dtype=np.int32, dims=(None, 100)), 49 | ] 50 | ) 51 | 52 | candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32) 53 | movie_ids_1 = np.zeros(100, dtype=np.int32) 54 | movie_ids_1[:20] = np.unique(candidate_ids)[:20] 55 | 56 | combined_features = { 57 | "candidate_ids": np.expand_dims(candidate_ids, axis=0), 58 | "movie_ids": np.expand_dims(movie_ids_1, axis=0), 59 | } 60 | 61 | request_data = TensorTable(combined_features) 62 | 63 | filtering = ["candidate_ids"] >> FilterCandidates(filter_out=["movie_ids"]) 64 | 65 | ensemble = Ensemble(filtering, request_schema) 66 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime, name=model_name) 67 | 68 | assert ensemble_config.name == expected_model_name 69 | response = run_ensemble_on_tritonserver( 70 | tmpdir, 71 | ensemble.input_schema, 72 | request_data, 73 | ensemble.output_schema.column_names, 74 | ensemble_config.name, 75 | ) 76 | assert response is not None 77 | assert len(response["filtered_ids"]) == 80 78 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/test_ensemble.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import random 17 | 18 | import numpy as np 19 | 20 | from merlin.dag.executors import LocalExecutor 21 | from merlin.schema import ColumnSchema, Schema 22 | from merlin.systems.dag.ensemble import Ensemble 23 | from merlin.systems.dag.ops.session_filter import FilterCandidates 24 | from merlin.table import TensorTable 25 | 26 | 27 | def test_ensemble_save_load(tmpdir): 28 | request_schema = Schema( 29 | [ 30 | ColumnSchema("candidate_ids", dtype=np.int32), 31 | ColumnSchema("movie_ids", dtype=np.int32), 32 | ] 33 | ) 34 | 35 | candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32) 36 | movie_ids_1 = np.zeros(100, dtype=np.int32) 37 | movie_ids_1[:20] = np.unique(candidate_ids)[:20] 38 | 39 | combined_features = { 40 | "candidate_ids": candidate_ids, 41 | "movie_ids": movie_ids_1, 42 | } 43 | 44 | request_data = TensorTable(combined_features) 45 | 46 | filtering = ["candidate_ids"] >> FilterCandidates(filter_out=["movie_ids"]) 47 | 48 | ensemble = Ensemble(filtering, request_schema) 49 | ensemble.save(str(tmpdir)) 50 | 51 | loaded_ensemble = Ensemble.load(str(tmpdir)) 52 | 53 | executor = LocalExecutor() 54 | response = executor.transform(request_data, [ensemble.graph.output_node]) 55 | 56 | loaded_response = executor.transform(request_data, [loaded_ensemble.graph.output_node]) 57 | 58 | assert loaded_response["filtered_ids"] == response["filtered_ids"] 59 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/test_executors.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import random 17 | import shutil 18 | 19 | import numpy as np 20 | import pandas as pd 21 | import pytest 22 | 23 | from merlin.core.dispatch import HAS_GPU, make_df 24 | from merlin.dag.executors import DaskExecutor, LocalExecutor 25 | from merlin.io import Dataset 26 | from merlin.schema import ColumnSchema, Schema 27 | from merlin.systems.dag.ensemble import Ensemble 28 | from merlin.systems.dag.ops.session_filter import FilterCandidates 29 | from merlin.table import TensorTable 30 | 31 | TRITON_SERVER_PATH = shutil.which("tritonserver") 32 | 33 | 34 | def test_run_dag_on_tensor_table_with_local_executor(): 35 | request_schema = Schema( 36 | [ 37 | ColumnSchema("candidate_ids", dtype=np.int32), 38 | ColumnSchema("movie_ids", dtype=np.int32), 39 | ] 40 | ) 41 | 42 | candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32) 43 | movie_ids_1 = np.zeros(100, dtype=np.int32) 44 | movie_ids_1[:20] = np.unique(candidate_ids)[:20] 45 | 46 | combined_features = { 47 | "candidate_ids": candidate_ids, 48 | "movie_ids": movie_ids_1, 49 | } 50 | 51 | request_data = TensorTable(combined_features) 52 | 53 | filtering = ["candidate_ids"] >> FilterCandidates(filter_out=["movie_ids"]) 54 | 55 | ensemble = Ensemble(filtering, request_schema) 56 | 57 | executor = LocalExecutor() 58 | response = executor.transform(request_data, [ensemble.graph.output_node]) 59 | 60 | assert response is not None 61 | assert isinstance(response, TensorTable) 62 | assert len(response["filtered_ids"]) == 80 63 | 64 | 65 | @pytest.mark.skipif(not HAS_GPU, reason="unable to find GPU") 66 | def test_run_dag_on_dataframe_with_local_executor(): 67 | import cudf 68 | 69 | request_schema = Schema( 70 | [ 71 | ColumnSchema("candidate_ids", dtype=np.int32), 72 | ColumnSchema("movie_ids", dtype=np.int32), 73 | ] 74 | ) 75 | 76 | candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32) 77 | movie_ids_1 = np.zeros(100, dtype=np.int32) 78 | movie_ids_1[:20] = np.unique(candidate_ids)[:20] 79 | 80 | combined_features = { 81 | "candidate_ids": candidate_ids, 82 | "movie_ids": movie_ids_1, 83 | } 84 | 85 | request_data = make_df(combined_features) 86 | 87 | filtering = ["candidate_ids"] >> FilterCandidates(filter_out=["movie_ids"]) 88 | ensemble = Ensemble(filtering, request_schema) 89 | 90 | executor = LocalExecutor() 91 | response = executor.transform(request_data, [ensemble.graph.output_node]) 92 | 93 | assert response is not None 94 | assert isinstance(response, (cudf.DataFrame, pd.DataFrame)) 95 | assert len(response["filtered_ids"]) == 80 96 | 97 | 98 | @pytest.mark.skipif(not HAS_GPU, reason="unable to find GPU") 99 | def test_run_dag_on_dataframe_with_dask_executor(): 100 | import dask_cudf 101 | 102 | request_schema = Schema( 103 | [ 104 | ColumnSchema("candidate_ids", dtype=np.int32), 105 | ColumnSchema("movie_ids", dtype=np.int32), 106 | ] 107 | ) 108 | 109 | candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32) 110 | movie_ids_1 = np.zeros(100, dtype=np.int32) 111 | movie_ids_1[:20] = np.unique(candidate_ids)[:20] 112 | 113 | combined_features = { 114 | "candidate_ids": candidate_ids, 115 | "movie_ids": movie_ids_1, 116 | } 117 | 118 | request_data = make_df(combined_features) 119 | request_dataset = Dataset(request_data) 120 | 121 | filtering = ["candidate_ids"] >> FilterCandidates(filter_out=["movie_ids"]) 122 | ensemble = Ensemble(filtering, request_schema) 123 | 124 | executor = DaskExecutor() 125 | response = executor.transform(request_dataset.to_ddf(), [ensemble.graph.output_node]) 126 | 127 | assert response is not None 128 | assert isinstance(response, dask_cudf.DataFrame) 129 | assert len(response["filtered_ids"]) == 80 130 | -------------------------------------------------------------------------------- /tests/unit/systems/dag/test_graph.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import pytest 17 | 18 | from merlin.dag.node import postorder_iter_nodes 19 | from merlin.dag.ops.concat_columns import ConcatColumns 20 | from merlin.dag.ops.selection import SelectionOp 21 | from merlin.schema import Schema 22 | from nvtabular import Workflow 23 | from nvtabular import ops as wf_ops 24 | 25 | ensemble = pytest.importorskip("merlin.systems.dag.ensemble") 26 | workflow_op = pytest.importorskip("merlin.systems.dag.ops.workflow") 27 | 28 | from merlin.systems.dag.ops.workflow import TransformWorkflow # noqa 29 | 30 | 31 | def test_inference_schema_propagation(): 32 | input_columns = ["a", "b", "c"] 33 | request_schema = Schema(input_columns) 34 | expected_schema = Schema(["a_nvt", "b_nvt", "c_nvt"]) 35 | 36 | # NVT 37 | workflow_ops = input_columns >> wf_ops.Rename(postfix="_nvt") 38 | workflow = Workflow(workflow_ops) 39 | workflow.fit_schema(request_schema) 40 | 41 | assert workflow.input_schema == request_schema 42 | assert workflow.output_schema == expected_schema 43 | 44 | # Triton 45 | triton_ops = input_columns >> workflow_op.TransformWorkflow(workflow) 46 | ensemble_out = ensemble.Ensemble(triton_ops, request_schema) 47 | 48 | assert ensemble_out.input_schema == request_schema 49 | assert ensemble_out.output_schema == expected_schema 50 | 51 | 52 | def test_graph_traverse_algo(): 53 | chain_1 = ["name-cat"] >> TransformWorkflow(Workflow(["name-cat"] >> wf_ops.Categorify())) 54 | chain_2 = ["name-string"] >> TransformWorkflow(Workflow(["name-string"] >> wf_ops.Categorify())) 55 | 56 | triton_chain = chain_1 + chain_2 57 | 58 | ordered_list = list(postorder_iter_nodes(triton_chain)) 59 | assert len(ordered_list) == 5 60 | assert isinstance(ordered_list[0].op, SelectionOp) 61 | assert isinstance(ordered_list[-1].op, ConcatColumns) 62 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/embedding_op.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2023, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | from merlin.core.compat import numpy as np 18 | from merlin.dataloader.ops.embeddings import NumpyEmbeddingOperator 19 | from merlin.schema import ColumnSchema, Schema, Tags 20 | from merlin.systems.dag.ensemble import Ensemble 21 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver 22 | from merlin.table import TensorTable 23 | 24 | 25 | def test_embedding_op_no_triton(tmpdir): 26 | embeddings = np.random.rand(100, 50) 27 | schema = Schema( 28 | [ColumnSchema("id", dtype=np.int32).with_tags([Tags.CATEGORICAL, Tags.EMBEDDING])] 29 | ) 30 | 31 | graph = ["id"] >> NumpyEmbeddingOperator(embeddings) 32 | triton_ens = Ensemble(graph, schema) 33 | req_table = TensorTable({"id": np.array([1, 2, 3])}) 34 | result = triton_ens.transform(req_table) 35 | assert ["id", "embeddings"] == result.columns 36 | assert result["embeddings"].shape.as_tuple == (3, 50) 37 | 38 | 39 | def test_embedding_op_triton(tmpdir): 40 | embeddings = np.random.rand(100, 50) 41 | schema = Schema( 42 | [ColumnSchema("id", dtype=np.int32).with_tags([Tags.CATEGORICAL, Tags.EMBEDDING])] 43 | ) 44 | 45 | graph = ["id"] >> NumpyEmbeddingOperator(embeddings) 46 | triton_ens = Ensemble(graph, schema) 47 | 48 | ensemble_config, node_configs = triton_ens.export(str(tmpdir)) 49 | 50 | req_table = TensorTable({"id": np.array([1, 2, 3], dtype=np.int32)}) 51 | response = run_ensemble_on_tritonserver( 52 | str(tmpdir), schema, req_table, ["id", "embeddings"], ensemble_config.name 53 | ) 54 | 55 | assert "embeddings" in response 56 | assert response["embeddings"].shape == (3, 50) 57 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/faiss/test_executor.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import shutil 17 | 18 | import numpy as np 19 | import pytest 20 | 21 | from merlin.core.dispatch import make_df 22 | from merlin.schema import ColumnSchema, Schema 23 | from merlin.systems.dag.ensemble import Ensemble 24 | from merlin.systems.dag.ops.faiss import QueryFaiss, setup_faiss 25 | 26 | TRITON_SERVER_PATH = shutil.which("tritonserver") 27 | pytest.importorskip("merlin.dataloader.tf_utils") 28 | from merlin.dataloader.tf_utils import configure_tensorflow # noqa 29 | 30 | tritonclient = pytest.importorskip("tritonclient") 31 | grpcclient = pytest.importorskip("tritonclient.grpc") 32 | 33 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver # noqa 34 | 35 | configure_tensorflow() 36 | 37 | import tensorflow as tf # noqa 38 | 39 | from merlin.systems.dag.ops.tensorflow import PredictTensorflow # noqa 40 | from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime # noqa 41 | from merlin.table import TensorTable # noqa 42 | 43 | 44 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 45 | def test_faiss_in_triton_executor_model(tmpdir): 46 | # Simulate a user vector with a TF model 47 | model = tf.keras.models.Sequential( 48 | [ 49 | tf.keras.Input(name="user_id", dtype=tf.int32, shape=(1,)), 50 | tf.keras.layers.Dense(128, activation="relu", name="output"), 51 | ] 52 | ) 53 | 54 | model.compile( 55 | optimizer="adam", 56 | loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True), 57 | metrics=[tf.metrics.SparseCategoricalAccuracy()], 58 | ) 59 | 60 | faiss_path = tmpdir / "faiss.index" 61 | item_ids = np.arange(0, 100) 62 | item_embeddings = np.random.rand(100, 128) 63 | # cannot turn a list column in cudf directly to numpy so must delegate to pandas as bridge 64 | df = make_df({"item_id": item_ids, "embedding": item_embeddings.tolist()}, device="cpu") 65 | setup_faiss(df, faiss_path) 66 | 67 | request_schema = Schema( 68 | [ 69 | ColumnSchema("user_id", dtype=np.int32, dims=(None, 1)), 70 | ] 71 | ) 72 | 73 | request_data = TensorTable( 74 | { 75 | "user_id": np.array([[1]], dtype=np.int32), 76 | } 77 | ) 78 | 79 | retrieval = ["user_id"] >> PredictTensorflow(model) >> QueryFaiss(faiss_path) 80 | 81 | ensemble = Ensemble(retrieval, request_schema) 82 | ensemble_config, _ = ensemble.export(tmpdir, runtime=TritonExecutorRuntime()) 83 | 84 | response = run_ensemble_on_tritonserver( 85 | tmpdir, 86 | ensemble.input_schema, 87 | request_data, 88 | ensemble.output_schema.column_names, 89 | ensemble_config.name, 90 | ) 91 | 92 | assert response is not None 93 | assert len(response["candidate_ids"]) == 10 94 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/feast/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | pytest.importorskip("feast") 20 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/fil/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | pytest.importorskip("xgboost") 20 | pytest.importorskip("sklearn") 21 | pytest.importorskip("lightgbm") 22 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/fil/test_ensemble.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import shutil 17 | 18 | import numpy as np 19 | import pandas as pd 20 | import pytest 21 | import sklearn.datasets 22 | import xgboost 23 | 24 | from merlin.dag import ColumnSelector 25 | from merlin.io import Dataset 26 | from merlin.schema import ColumnSchema, Schema 27 | from merlin.systems.dag.ops.fil import PredictForest 28 | from nvtabular import Workflow 29 | from nvtabular import ops as wf_ops 30 | 31 | triton = pytest.importorskip("merlin.systems.triton") 32 | export = pytest.importorskip("merlin.systems.dag.ensemble") 33 | 34 | from merlin.systems.dag.ensemble import Ensemble # noqa 35 | from merlin.systems.dag.ops.workflow import TransformWorkflow # noqa 36 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver # noqa 37 | 38 | TRITON_SERVER_PATH = shutil.which("tritonserver") 39 | 40 | 41 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 42 | def test_workflow_with_forest_inference(tmpdir): 43 | rows = 200 44 | num_features = 16 45 | X, y = sklearn.datasets.make_regression( 46 | n_samples=rows, 47 | n_features=num_features, 48 | n_informative=num_features // 3, 49 | random_state=0, 50 | ) 51 | feature_names = [str(i) for i in range(num_features)] 52 | df = pd.DataFrame(X, columns=feature_names, dtype=np.float32) 53 | dataset = Dataset(df) 54 | 55 | # Fit GBDT Model 56 | model = xgboost.XGBRegressor() 57 | model.fit(X, y) 58 | 59 | input_column_schemas = [ColumnSchema(col, dtype=np.float32) for col in feature_names] 60 | input_schema = Schema(input_column_schemas) 61 | selector = ColumnSelector(feature_names) 62 | 63 | workflow_ops = feature_names >> wf_ops.LogOp() 64 | workflow = Workflow(workflow_ops) 65 | workflow.fit(dataset) 66 | 67 | triton_chain = selector >> TransformWorkflow(workflow) >> PredictForest(model, input_schema) 68 | 69 | triton_ens = Ensemble(triton_chain, input_schema) 70 | 71 | request_df = df[:5] 72 | ensemble_config, _ = triton_ens.export(tmpdir) 73 | 74 | response = run_ensemble_on_tritonserver( 75 | str(tmpdir), input_schema, request_df, ["output__0"], ensemble_config.name 76 | ) 77 | assert response["output__0"].shape == (5,) 78 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/implicit/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | pytest.importorskip("implicit") 20 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/implicit/test_executor.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import shutil 17 | 18 | import implicit 19 | import numpy as np 20 | import pytest 21 | from scipy.sparse import csr_matrix 22 | 23 | from merlin.core.dispatch import make_df 24 | from merlin.schema import ColumnSchema, Schema 25 | from merlin.systems.dag.ensemble import Ensemble 26 | from merlin.systems.dag.ops.implicit import PredictImplicit 27 | 28 | TRITON_SERVER_PATH = shutil.which("tritonserver") 29 | 30 | tritonclient = pytest.importorskip("tritonclient") 31 | grpcclient = pytest.importorskip("tritonclient.grpc") 32 | 33 | from merlin.systems.dag.runtimes.triton import TritonExecutorRuntime # noqa 34 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver # noqa 35 | 36 | 37 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 38 | @pytest.mark.parametrize("runtime", [None, TritonExecutorRuntime()]) 39 | def test_implicit_in_triton_executor_model(tmpdir, runtime): 40 | model = implicit.bpr.BayesianPersonalizedRanking() 41 | n = 100 42 | user_items = csr_matrix(np.random.choice([0, 1], size=n * n, p=[0.9, 0.1]).reshape(n, n)) 43 | model.fit(user_items) 44 | 45 | request_schema = Schema([ColumnSchema("user_id", dtype="int64")]) 46 | 47 | implicit_op = PredictImplicit(model, num_to_recommend=10) 48 | triton_chain = request_schema.column_names >> implicit_op 49 | 50 | ensemble = Ensemble(triton_chain, request_schema) 51 | ensemble_config, _ = ensemble.export(tmpdir, runtime=runtime) 52 | 53 | input_user_id = np.array([0, 1], dtype=np.int64) 54 | 55 | response = run_ensemble_on_tritonserver( 56 | tmpdir, 57 | request_schema, 58 | make_df({"user_id": input_user_id}), 59 | ensemble.output_schema.column_names, 60 | ensemble_config.name, 61 | ) 62 | assert response is not None 63 | assert len(response["ids"]) == len(input_user_id) 64 | assert len(response["scores"]) == len(input_user_id) 65 | assert len(response["ids"][0]) == 10 66 | assert len(response["scores"][0]) == 10 67 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/implicit/test_op.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import shutil 18 | 19 | import implicit 20 | import numpy as np 21 | import pytest 22 | from scipy.sparse import csr_matrix 23 | from tritonclient import grpc as grpcclient 24 | 25 | from merlin.schema import ColumnSchema, Schema 26 | from merlin.systems.dag.ensemble import Ensemble 27 | from merlin.systems.dag.ops.implicit import PredictImplicit 28 | from merlin.systems.triton.utils import run_triton_server 29 | 30 | TRITON_SERVER_PATH = shutil.which("tritonserver") 31 | 32 | 33 | triton = pytest.importorskip("merlin.systems.triton") 34 | 35 | 36 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 37 | @pytest.mark.parametrize( 38 | "model_cls", 39 | [ 40 | implicit.bpr.BayesianPersonalizedRanking, 41 | implicit.als.AlternatingLeastSquares, 42 | implicit.lmf.LogisticMatrixFactorization, 43 | ], 44 | ) 45 | def test_ensemble(model_cls, tmpdir): 46 | model = model_cls() 47 | n = 100 48 | user_items = csr_matrix(np.random.choice([0, 1], size=n * n, p=[0.9, 0.1]).reshape(n, n)) 49 | model.fit(user_items) 50 | 51 | num_to_recommend = np.random.randint(1, n) 52 | 53 | user_items = None 54 | ids, scores = model.recommend( 55 | [0, 1], user_items, N=num_to_recommend, filter_already_liked_items=False 56 | ) 57 | 58 | implicit_op = PredictImplicit(model, num_to_recommend=num_to_recommend) 59 | 60 | input_schema = Schema([ColumnSchema("user_id", dtype="int64", dims=(None, 1))]) 61 | 62 | triton_chain = input_schema.column_names >> implicit_op 63 | 64 | triton_ens = Ensemble(triton_chain, input_schema) 65 | ensemble_config, _ = triton_ens.export(tmpdir) 66 | 67 | input_user_id = np.array([[0], [1]], dtype=np.int64) 68 | inputs = [ 69 | grpcclient.InferInput( 70 | "user_id", input_user_id.shape, triton.np_to_triton_dtype(input_user_id.dtype) 71 | ), 72 | ] 73 | inputs[0].set_data_from_numpy(input_user_id) 74 | outputs = [grpcclient.InferRequestedOutput("scores"), grpcclient.InferRequestedOutput("ids")] 75 | 76 | response = None 77 | 78 | with run_triton_server(tmpdir) as client: 79 | response = client.infer(ensemble_config.name, inputs, outputs=outputs) 80 | 81 | response_ids = response.as_numpy("ids") 82 | response_scores = response.as_numpy("scores") 83 | 84 | np.testing.assert_array_equal(ids, response_ids) 85 | np.testing.assert_array_equal(scores, response_scores) 86 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/nvtabular/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | pytest.importorskip("nvtabular") 20 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/padding_op.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2023, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | from merlin.core.compat import numpy as np 18 | from merlin.dataloader.ops.padding import Padding 19 | from merlin.schema import ColumnSchema, Schema 20 | from merlin.systems.dag.ensemble import Ensemble 21 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver 22 | from merlin.table import TensorTable 23 | 24 | 25 | def test_padding_op_no_triton(tmpdir): 26 | padding_size = 5 27 | padding_value = 0 28 | req_table = TensorTable( 29 | {"a": (np.array([1, 2, 3], dtype=np.int32), np.array([0, 1, 3], dtype=np.int32))} 30 | ) 31 | schema = Schema( 32 | [ColumnSchema("a", dtype=np.int32, dims=req_table["a"].shape, is_list=True, is_ragged=True)] 33 | ) 34 | 35 | graph = ["a"] >> Padding(padding_size, padding_value) 36 | triton_ens = Ensemble(graph, schema) 37 | result = triton_ens.transform(req_table) 38 | 39 | assert ["a"] == result.columns 40 | assert result["a"].values.shape == (2, padding_size) 41 | 42 | 43 | def test_padding_op_triton(tmpdir): 44 | padding_size = 5 45 | padding_value = 0 46 | req_table = TensorTable( 47 | {"a": (np.array([1, 2, 3], dtype=np.int32), np.array([0, 1, 3], dtype=np.int32))} 48 | ) 49 | schema = Schema( 50 | [ColumnSchema("a", dtype=np.int32, dims=req_table["a"].shape, is_list=True, is_ragged=True)] 51 | ) 52 | 53 | graph = ["a"] >> Padding(padding_size, padding_value) 54 | 55 | triton_ens = Ensemble(graph, schema) 56 | ensemble_config, node_configs = triton_ens.export(str(tmpdir)) 57 | response = run_ensemble_on_tritonserver( 58 | str(tmpdir), schema, req_table, ["a"], ensemble_config.name 59 | ) 60 | 61 | assert "a" in response 62 | assert response["a"].shape == (2, padding_size) 63 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/tf/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | pytest.importorskip("tensorflow") 20 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/torch/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | pytest.importorskip("torch") 20 | -------------------------------------------------------------------------------- /tests/unit/systems/ops/torch/test_ensemble.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2023, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import shutil 18 | 19 | import numpy as np 20 | import pandas as pd 21 | import pytest 22 | import tritonclient.utils 23 | 24 | from merlin.schema import ColumnSchema, Schema 25 | from merlin.systems.dag.ensemble import Ensemble 26 | from merlin.systems.dag.ops.pytorch import PredictPyTorch 27 | from merlin.systems.triton.utils import run_ensemble_on_tritonserver 28 | 29 | torch = pytest.importorskip("torch") 30 | 31 | TRITON_SERVER_PATH = shutil.which("tritonserver") 32 | 33 | 34 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 35 | def test_model_in_ensemble(tmpdir): 36 | class MyModel(torch.nn.Module): 37 | def forward(self, x): 38 | v = torch.stack(list(x.values())).sum(axis=0) 39 | return v 40 | 41 | model = MyModel() 42 | 43 | traced_model = torch.jit.trace(model, {"a": torch.tensor(1), "b": torch.tensor(2)}, strict=True) 44 | 45 | model_input_schema = Schema( 46 | [ColumnSchema("a", dtype="int64"), ColumnSchema("b", dtype="int64")] 47 | ) 48 | model_output_schema = Schema([ColumnSchema("output", dtype="int64")]) 49 | 50 | model_node = model_input_schema.column_names >> PredictPyTorch( 51 | traced_model, model_input_schema, model_output_schema 52 | ) 53 | 54 | ensemble = Ensemble(model_node, model_input_schema) 55 | 56 | ensemble_config, _ = ensemble.export(str(tmpdir)) 57 | 58 | df = pd.DataFrame({"a": [1], "b": [2]}) 59 | 60 | response = run_ensemble_on_tritonserver( 61 | str(tmpdir), model_input_schema, df, ["output"], ensemble_config.name 62 | ) 63 | np.testing.assert_array_equal(response["output"], np.array([3])) 64 | 65 | 66 | @pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found") 67 | def test_model_error(tmpdir): 68 | class MyModel(torch.nn.Module): 69 | def forward(self, x): 70 | v = torch.stack(list(x.values())).sum() 71 | return v 72 | 73 | model = MyModel() 74 | 75 | traced_model = torch.jit.trace(model, {"a": torch.tensor(1), "b": torch.tensor(2)}, strict=True) 76 | 77 | model_input_schema = Schema([ColumnSchema("a", dtype="int64")]) 78 | model_output_schema = Schema([ColumnSchema("output", dtype="int64")]) 79 | 80 | model_node = model_input_schema.column_names >> PredictPyTorch( 81 | traced_model, model_input_schema, model_output_schema 82 | ) 83 | 84 | ensemble = Ensemble(model_node, model_input_schema) 85 | 86 | ensemble_config, _ = ensemble.export(str(tmpdir)) 87 | 88 | # run inference with missing input (that was present when model was compiled) 89 | # we're expecting a KeyError at runtime. 90 | df = pd.DataFrame({"a": [1]}) 91 | 92 | with pytest.raises(tritonclient.utils.InferenceServerException) as exc_info: 93 | run_ensemble_on_tritonserver( 94 | str(tmpdir), model_input_schema, df, ["output"], ensemble_config.name 95 | ) 96 | assert "The following operation failed in the TorchScript interpreter" in str(exc_info.value) 97 | assert "RuntimeError: KeyError: b" in str(exc_info.value) 98 | -------------------------------------------------------------------------------- /tests/unit/systems/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /tests/unit/systems/utils/ops.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | import pytest 18 | 19 | from merlin.core.protocols import Transformable 20 | from merlin.dag import ColumnSelector 21 | from merlin.table import TensorTable 22 | 23 | inf_op = pytest.importorskip("merlin.systems.dag.runtimes.triton.ops.operator") 24 | 25 | 26 | class PlusTwoOp(inf_op.TritonOperator): 27 | def __init__(self): 28 | super().__init__(self) 29 | 30 | def transform( 31 | self, col_selector: ColumnSelector, transformable: Transformable 32 | ) -> Transformable: 33 | result = TensorTable() 34 | 35 | for name, column in transformable.items(): 36 | result[f"{name}_plus_2"] = type(column)(column.values + 2, column.offsets) 37 | 38 | return result 39 | 40 | def column_mapping(self, col_selector): 41 | column_mapping = {} 42 | for col_name in col_selector.names: 43 | column_mapping[f"{col_name}_plus_2"] = [col_name] 44 | return column_mapping 45 | 46 | @classmethod 47 | def from_config(cls, config, **kwargs): 48 | return PlusTwoOp() 49 | -------------------------------------------------------------------------------- /tests/unit/systems/utils/tf.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import pytest 17 | 18 | loader_tf_utils = pytest.importorskip("nvtabular.loader.tf_utils") # noqa 19 | loader_tf_utils.configure_tensorflow() 20 | tf = pytest.importorskip("tensorflow") 21 | 22 | from nvtabular.framework_utils.tensorflow import layers # noqa 23 | 24 | 25 | class ExampleModel(tf.keras.Model): 26 | def __init__(self, cat_columns, cat_mh_columns, embed_tbl_shapes): 27 | super(ExampleModel, self).__init__() 28 | 29 | self.cat_columns = cat_columns 30 | self.cat_mh_columns = cat_mh_columns 31 | self.embed_tbl_shapes = embed_tbl_shapes 32 | 33 | self.emb_layers = [] # output of all embedding layers, which will be concatenated 34 | for col in self.cat_columns + self.cat_mh_columns: 35 | self.emb_layers.append( 36 | tf.feature_column.embedding_column( 37 | tf.feature_column.categorical_column_with_identity( 38 | col, self.embed_tbl_shapes[col][0] 39 | ), # Input dimension (vocab size) 40 | self.embed_tbl_shapes[col][1], # Embedding output dimension 41 | ) 42 | ) 43 | self.emb_layer = layers.DenseFeatures(self.emb_layers) 44 | self.dense_layer = tf.keras.layers.Dense(128, activation="relu") 45 | self.output_layer = tf.keras.layers.Dense(1, activation="sigmoid", name="output") 46 | 47 | def call(self, inputs): 48 | reshaped_inputs = {} 49 | for input_name, input_ in inputs.items(): 50 | reshaped_inputs[input_name] = tf.reshape(input_, (-1, 1)) 51 | 52 | x_emb_output = self.emb_layer(reshaped_inputs) 53 | x = self.dense_layer(x_emb_output) 54 | x = self.output_layer(x) 55 | return {"predictions": x} 56 | 57 | 58 | def create_tf_model(cat_columns: list, cat_mh_columns: list, embed_tbl_shapes: dict): 59 | 60 | model = ExampleModel(cat_columns, cat_columns, embed_tbl_shapes) 61 | model.compile("sgd", "binary_crossentropy") 62 | example_input = tf.constant([1, 2, 3], dtype=tf.int64) 63 | cat_column_data = {cat_col: example_input for cat_col in cat_columns} 64 | model(cat_column_data) 65 | return model 66 | -------------------------------------------------------------------------------- /tests/unit/systems/utils/torch.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import pytest 17 | 18 | torch = pytest.importorskip("torch") # noqa 19 | 20 | from nvtabular.framework_utils.torch.models import Model # noqa 21 | 22 | 23 | def create_pytorch_model(cat_columns: list, cat_mh_columns: list, embed_tbl_shapes: dict): 24 | single_hot = {k: v for k, v in embed_tbl_shapes.items() if k in cat_columns} 25 | multi_hot = {k: v for k, v in embed_tbl_shapes.items() if k in cat_mh_columns} 26 | model = Model( 27 | embedding_table_shapes=(single_hot, multi_hot), 28 | num_continuous=0, 29 | emb_dropout=0.0, 30 | layer_hidden_dims=[128, 128, 128], 31 | layer_dropout_rates=[0.0, 0.0, 0.0], 32 | ).to("cuda") 33 | return model 34 | -------------------------------------------------------------------------------- /tests/version/test_version.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2022, NVIDIA CORPORATION. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | import pytest 17 | from packaging.version import Version 18 | 19 | import merlin.systems 20 | 21 | 22 | @pytest.mark.version 23 | def test_version(): 24 | """test to get back version of library""" 25 | assert Version(merlin.systems.__version__) >= Version("0.6.0") 26 | --------------------------------------------------------------------------------