├── .github ├── ISSUE_TEMPLATE │ ├── bug-report.yml │ ├── config.yml │ ├── documentation.yml │ └── feature-request.yml └── workflows │ ├── glt-ci.yml │ ├── glt-v6d-ci.yml │ ├── manylinux-cd.yml │ └── scripts │ ├── build.sh │ ├── build_wheel.sh │ └── upload_pypi.sh ├── .gitignore ├── .gitmodules ├── .readthedocs.yaml ├── CMakeLists.txt ├── LICENSE ├── README.md ├── benchmarks └── api │ ├── README.md │ ├── bench_dist_config.yml │ ├── bench_dist_neighbor_loader.py │ ├── bench_feature.py │ ├── bench_sampler.py │ └── run_dist_bench.py ├── dockerfiles ├── graphlearn-torch-dev.Dockerfile ├── graphlearn-torch-wheel.Dockerfile └── gs-graphlearn-torch-dev.Dockerfile ├── docs ├── Makefile ├── README.md ├── apis │ ├── channel.rst │ ├── data.rst │ ├── distributed.rst │ ├── loader.rst │ ├── modules.rst │ ├── partition.rst │ ├── sampler.rst │ └── utils.rst ├── conf.py ├── contribute.md ├── faq.md ├── figures │ ├── arch.png │ ├── dist_arch_server_client_mode.png │ ├── dist_arch_worker_mode.png │ ├── glt_dgl.png │ ├── replica_per_dg.png │ ├── replica_per_gpu.png │ ├── scale_out.png │ ├── scale_up.png │ └── uni_tensor.png ├── get_started │ ├── dist_train.md │ ├── graph_basic.md │ ├── hetero.md │ ├── link_pred.md │ └── node_class.md ├── index.rst ├── install │ └── install.md ├── make.bat ├── requirements.txt └── tutorial │ ├── aliyun.md │ ├── basic_object.md │ ├── dist.md │ └── graph_ops.md ├── examples ├── README.md ├── distributed │ ├── README.md │ ├── dist_sage_unsup │ │ ├── dist_sage_unsup.py │ │ └── preprocess_template.py │ ├── dist_train_sage_sup_config.yml │ ├── dist_train_sage_supervised.py │ ├── partition_ogbn_dataset.py │ ├── run_dist_train_sage_sup.py │ └── server_client_mode │ │ ├── sage_supervised_client.py │ │ └── sage_supervised_server.py ├── feature_mp.py ├── gpt │ ├── README.md │ ├── arxiv.py │ └── utils.py ├── graph_sage_unsup_ppi.py ├── hetero │ ├── bipartite_sage_unsup.py │ ├── hierarchical_sage.py │ ├── train_hgt_mag.py │ └── train_hgt_mag_mp.py ├── igbh │ ├── Dockerfile │ ├── README.md │ ├── __init__.py │ ├── build_partition_feature.py │ ├── compress_graph.py │ ├── dataset.py │ ├── dist_train_rgnn.py │ ├── download.py │ ├── download_igbh_full.sh │ ├── mlperf_logging_utils.py │ ├── partition.py │ ├── rgnn.py │ ├── split_seeds.py │ ├── train_rgnn_multi_gpu.py │ └── utilities.py ├── multi_gpu │ └── train_sage_ogbn_papers100m.py ├── pai │ ├── README.md │ ├── ogbn_products │ │ ├── README.md │ │ ├── data_preprocess.py │ │ ├── dist_train_products_sage.py │ │ └── train_products_sage.py │ └── requirements.txt ├── seal_link_pred.py ├── train_sage_ogbn_products.py └── train_sage_prod_with_trim.py ├── graphlearn_torch ├── csrc │ ├── cpu │ │ ├── graph.cc │ │ ├── inducer.cc │ │ ├── inducer.h │ │ ├── random_negative_sampler.cc │ │ ├── random_negative_sampler.h │ │ ├── random_sampler.cc │ │ ├── random_sampler.h │ │ ├── stitch_sample_results.cc │ │ ├── subgraph_op.cc │ │ ├── subgraph_op.h │ │ ├── weighted_sampler.cc │ │ └── weighted_sampler.h │ ├── cuda │ │ ├── graph.cu │ │ ├── hash_table.cu │ │ ├── inducer.cu │ │ ├── inducer.cuh │ │ ├── random_negative_sampler.cu │ │ ├── random_negative_sampler.cuh │ │ ├── random_sampler.cu │ │ ├── random_sampler.cuh │ │ ├── stitch_sample_results.cu │ │ ├── subgraph_op.cu │ │ ├── subgraph_op.cuh │ │ ├── unified_tensor.cu │ │ └── weighted_sampler.cuh │ ├── sample_queue.cc │ ├── shm_queue.cc │ └── tensor_map.cc ├── include │ ├── common.cuh │ ├── common.h │ ├── func_factory.h │ ├── graph.h │ ├── hash_table.cuh │ ├── inducer_base.h │ ├── negative_sampler.h │ ├── sample_queue.h │ ├── sampler.h │ ├── shm_queue.h │ ├── stitch_sample_results.h │ ├── subgraph_op_base.h │ ├── tensor_map.h │ ├── types.h │ └── unified_tensor.cuh ├── python │ ├── __init__.py │ ├── channel │ │ ├── __init__.py │ │ ├── base.py │ │ ├── mp_channel.py │ │ ├── remote_channel.py │ │ └── shm_channel.py │ ├── data │ │ ├── __init__.py │ │ ├── dataset.py │ │ ├── feature.py │ │ ├── graph.py │ │ ├── reorder.py │ │ ├── table_dataset.py │ │ ├── unified_tensor.py │ │ └── vineyard_utils.py │ ├── distributed │ │ ├── __init__.py │ │ ├── dist_client.py │ │ ├── dist_context.py │ │ ├── dist_dataset.py │ │ ├── dist_feature.py │ │ ├── dist_graph.py │ │ ├── dist_link_neighbor_loader.py │ │ ├── dist_loader.py │ │ ├── dist_neighbor_loader.py │ │ ├── dist_neighbor_sampler.py │ │ ├── dist_options.py │ │ ├── dist_random_partitioner.py │ │ ├── dist_sampling_producer.py │ │ ├── dist_server.py │ │ ├── dist_subgraph_loader.py │ │ ├── dist_table_dataset.py │ │ ├── event_loop.py │ │ └── rpc.py │ ├── loader │ │ ├── __init__.py │ │ ├── link_loader.py │ │ ├── link_neighbor_loader.py │ │ ├── neighbor_loader.py │ │ ├── node_loader.py │ │ ├── subgraph_loader.py │ │ └── transform.py │ ├── partition │ │ ├── __init__.py │ │ ├── base.py │ │ ├── frequency_partitioner.py │ │ ├── partition_book.py │ │ └── random_partitioner.py │ ├── py_export_glt.cc │ ├── py_export_v6d.cc │ ├── sampler │ │ ├── __init__.py │ │ ├── base.py │ │ ├── negative_sampler.py │ │ └── neighbor_sampler.py │ ├── typing.py │ └── utils │ │ ├── __init__.py │ │ ├── build_glt.py │ │ ├── common.py │ │ ├── device.py │ │ ├── exit_status.py │ │ ├── mixin.py │ │ ├── singleton.py │ │ ├── tensor.py │ │ ├── topo.py │ │ └── units.py └── v6d │ ├── vineyard_utils.cc │ └── vineyard_utils.h ├── install_dependencies.sh ├── scripts ├── build_wheel.sh ├── launch_container.sh ├── remove_container.sh ├── run_cpp_ut.sh └── run_python_ut.sh ├── setup.py ├── test ├── cpp │ ├── test_graph.cu │ ├── test_hash_table.cu │ ├── test_hetero_inducer.cu │ ├── test_inducer.cu │ ├── test_random_negative_sampler.cu │ ├── test_random_sampler.cu │ ├── test_shm_queue.cu │ ├── test_stitch_sample_results.cu │ ├── test_subgraph.cu │ ├── test_tensor_map_serializer.cu │ ├── test_unified_tensor.cu │ ├── test_vineyard.cc │ └── test_vineyard.cu └── python │ ├── dist_test_utils.py │ ├── test_dist_feature.py │ ├── test_dist_link_loader.py │ ├── test_dist_neighbor_loader.py │ ├── test_dist_random_partitioner.py │ ├── test_dist_subgraph_loader.py │ ├── test_feature.py │ ├── test_graph.py │ ├── test_hetero_neighbor_sampler.py │ ├── test_link_loader.py │ ├── test_neighbor_sampler.py │ ├── test_partition.py │ ├── test_pyg_remote_backend.py │ ├── test_sample_prob.py │ ├── test_shm_channel.py │ ├── test_subgraph.py │ ├── test_unified_tensor.py │ ├── test_vineyard.py │ └── vineyard_data │ └── modern_graph │ ├── created.csv │ ├── knows.csv │ ├── person.csv │ └── software.csv └── third_party └── googletest └── build.sh /.github/ISSUE_TEMPLATE/bug-report.yml: -------------------------------------------------------------------------------- 1 | name: 🐛 Bug Report 2 | description: Create a report to help us reproduce and fix the bug 3 | 4 | body: 5 | - type: markdown 6 | attributes: 7 | value: > 8 | #### Before submitting a bug, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/alibaba/graphlearn-for-pytorch/issues). 9 | - type: textarea 10 | attributes: 11 | label: 🐛 Describe the bug 12 | description: | 13 | Please provide a clear and concise description of what the bug is. 14 | 15 | If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. 16 | 17 | If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com. 18 | 19 | Please also paste or describe the results you observe instead of the expected results. If you observe an error, please paste the error message including the **full** traceback of the exception. It may be relevant to wrap error messages in ```` ```triple quotes blocks``` ````. 20 | placeholder: | 21 | A clear and concise description of what the bug is. 22 | 23 | ```python 24 | # Sample code to reproduce the problem 25 | ``` 26 | 27 | ``` 28 | The error message you got, with the full traceback. 29 | ``` 30 | validations: 31 | required: true 32 | 33 | - type: textarea 34 | attributes: 35 | label: Environment 36 | description: | 37 | Please provide as much information as possible about your environment. 38 | value: | 39 | * GLT version: 40 | * PyG version: 41 | * PyTorch version: 42 | * OS: 43 | * Python version: 44 | * CUDA/cuDNN version: 45 | * Any other relevant information 46 | - type: markdown 47 | attributes: 48 | value: > 49 | Thanks for contributing 🎉! 50 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/config.yml: -------------------------------------------------------------------------------- 1 | blank_issues_enabled: false 2 | contact_links: 3 | - name: 💬 Ask a Question 4 | url: https://github.com/alibaba/graphlearn-for-pytorch/discussions 5 | about: Ask and answer GLT related questions 6 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/documentation.yml: -------------------------------------------------------------------------------- 1 | name: 📚 Documentation 2 | description: Tell us about how we can improve our documentation https://graphlearn-torch.readthedocs.io/en/latest/ 3 | 4 | body: 5 | - type: textarea 6 | attributes: 7 | label: 📚 The doc issue 8 | description: > 9 | A clear and concise description of what content in https://graphlearn-torch.readthedocs.io/en/latest/ is an issue.. 10 | validations: 11 | required: true 12 | - type: textarea 13 | attributes: 14 | label: Suggest a potential alternative/fix 15 | description: > 16 | Tell us how we could improve the documentation in this regard. 17 | - type: markdown 18 | attributes: 19 | value: > 20 | Thanks for contributing 🎉! -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature-request.yml: -------------------------------------------------------------------------------- 1 | name: 🚀 Feature request 2 | description: Submit a proposal/request for a new GLT feature 3 | 4 | body: 5 | - type: textarea 6 | attributes: 7 | label: 🚀 The feature, motivation and pitch 8 | description: > 9 | A clear and concise description of the feature proposal. Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., *"I'm working on X and would like Y to be possible"*. If this is related to another GitHub issue, please link here too. 10 | validations: 11 | required: true 12 | - type: textarea 13 | attributes: 14 | label: Alternatives 15 | description: > 16 | A description of any alternative solutions or features you've considered, if any. 17 | - type: textarea 18 | attributes: 19 | label: Additional context 20 | description: > 21 | Add any other context or screenshots about the feature request. 22 | - type: markdown 23 | attributes: 24 | value: > 25 | Thanks for contributing 🎉! 26 | -------------------------------------------------------------------------------- /.github/workflows/glt-ci.yml: -------------------------------------------------------------------------------- 1 | name: GLT CI 2 | 3 | on: 4 | push: 5 | branches: [ main ] 6 | paths-ignore: 7 | - 'benchmarks/**' 8 | - 'docs/**' 9 | - 'examples/**' 10 | - '**.md' 11 | pull_request: 12 | branches: [ main ] 13 | paths-ignore: 14 | - 'benchmarks/**' 15 | - 'docs/**' 16 | - 'examples/**' 17 | - '**.md' 18 | 19 | env: 20 | IMAGE: graphlearn-torch-dev 21 | JOBNAME: glt-ci-${{ github.run_id }} 22 | DESTDIR: /mnt/graphlearn_for_pytorch 23 | 24 | jobs: 25 | run-glt-unittests: 26 | runs-on: self-hosted 27 | if: ${{ github.repository == 'alibaba/graphlearn-for-pytorch' }} 28 | steps: 29 | - name: Checkout Code 30 | uses: actions/checkout@v3 31 | 32 | - name: Update submodules 33 | working-directory: ${{github.workspace}} 34 | run: | 35 | git submodule update --init 36 | 37 | - name: Build image 38 | working-directory: ${{github.workspace}} 39 | run: | 40 | docker build -f dockerfiles/${IMAGE}.Dockerfile -t ${IMAGE} . 41 | 42 | - name: Launch CI container 43 | working-directory: ${{github.workspace}} 44 | run: | 45 | bash scripts/launch_container.sh ${IMAGE} ${JOBNAME} ${DESTDIR} 46 | 47 | - name: Build wheel 48 | run: | 49 | docker exec ${JOBNAME} bash ${DESTDIR}/scripts/build_wheel.sh 50 | 51 | - name: Run cpp unit tests 52 | run: | 53 | docker exec ${JOBNAME} bash ${DESTDIR}/scripts/run_cpp_ut.sh 54 | 55 | - name: Run python unit tests 56 | run: | 57 | docker exec ${JOBNAME} bash ${DESTDIR}/scripts/run_python_ut.sh 58 | 59 | - name: Remove CI container 60 | if: always() 61 | working-directory: ${{github.workspace}} 62 | run: | 63 | bash scripts/remove_container.sh ${JOBNAME} 64 | -------------------------------------------------------------------------------- /.github/workflows/glt-v6d-ci.yml: -------------------------------------------------------------------------------- 1 | name: GLT V6D CI 2 | 3 | on: 4 | push: 5 | branches: [ main ] 6 | paths-ignore: 7 | - 'benchmarks/**' 8 | - 'docs/**' 9 | - 'examples/**' 10 | - '**.md' 11 | pull_request: 12 | branches: [ main ] 13 | paths-ignore: 14 | - 'benchmarks/**' 15 | - 'docs/**' 16 | - 'examples/**' 17 | - '**.md' 18 | 19 | env: 20 | IMAGE: gs-graphlearn-torch-dev 21 | JOBNAME: glt-v6d-ci-${{ github.run_id }} 22 | DESTDIR: /home/graphscope/graphlearn-for-pytorch 23 | 24 | jobs: 25 | run-glt-v6d-unittests: 26 | runs-on: self-hosted 27 | if: ${{ github.repository == 'alibaba/graphlearn-for-pytorch' }} 28 | steps: 29 | - name: Checkout Code 30 | uses: actions/checkout@v3 31 | 32 | - name: Update submodules 33 | working-directory: ${{github.workspace}} 34 | run: | 35 | git submodule update --init 36 | 37 | - name: Build image 38 | working-directory: ${{github.workspace}} 39 | run: | 40 | docker build -f dockerfiles/${IMAGE}.Dockerfile -t ${IMAGE} . 41 | 42 | - name: Launch CI container 43 | working-directory: ${{github.workspace}} 44 | run: | 45 | WITH_VINEYARD=ON WITH_CUDA=OFF MOUNT_VOLUME=FALSE bash scripts/launch_container.sh ${IMAGE} ${JOBNAME} ${DESTDIR} && \ 46 | docker exec ${JOBNAME} sudo chown -R graphscope ${DESTDIR} && \ 47 | docker exec ${JOBNAME} sudo rm -rf ${DESTDIR}/dist ${DESTDIR}/build ${DESTDIR}/built 48 | 49 | - name: Build wheel 50 | run: | 51 | docker exec ${JOBNAME} bash ${DESTDIR}/scripts/build_wheel.sh 52 | 53 | - name: Run python unit tests 54 | run: | 55 | docker exec ${JOBNAME} python3 ${DESTDIR}/test/python/test_vineyard.py 56 | 57 | - name: Remove CI container 58 | if: always() 59 | working-directory: ${{github.workspace}} 60 | run: | 61 | bash scripts/remove_container.sh ${JOBNAME} -------------------------------------------------------------------------------- /.github/workflows/manylinux-cd.yml: -------------------------------------------------------------------------------- 1 | name: GLT CD 2 | 3 | on: 4 | push: 5 | tags: 6 | - '*' 7 | 8 | env: 9 | IMAGE: graphlearn-torch-wheel 10 | JOBNAME: glt-cd-${{ github.run_id }} 11 | DESTDIR: /mnt/graphlearn_for_pytorch 12 | 13 | jobs: 14 | build: 15 | runs-on: self-hosted 16 | if: ${{ github.repository == 'alibaba/graphlearn-for-pytorch' }} 17 | steps: 18 | - name: Checkout Code 19 | uses: actions/checkout@v3 20 | 21 | - name: Update submodules 22 | working-directory: ${{github.workspace}} 23 | run: | 24 | git submodule update --init 25 | 26 | - name: Build image 27 | working-directory: ${{github.workspace}} 28 | run: | 29 | docker build -f dockerfiles/${IMAGE}.Dockerfile -t ${IMAGE} . 30 | 31 | - name: Launch CD container 32 | working-directory: ${{github.workspace}} 33 | run: | 34 | bash scripts/launch_container.sh ${IMAGE} ${JOBNAME} ${DESTDIR} 35 | 36 | - name: Build python packages 37 | if: startsWith(github.ref, 'refs/tags/') 38 | shell: bash 39 | run: | 40 | docker exec ${JOBNAME} bash ${DESTDIR}/.github/workflows/scripts/build.sh ${GITHUB_REF} 41 | 42 | - name: Upload python packages 43 | if: startsWith(github.ref, 'refs/tags/') 44 | shell: bash 45 | env: 46 | PYPI_PWD: ${{ secrets.PYPI_PWD }} 47 | run: | 48 | bash .github/workflows/scripts/upload_pypi.sh 49 | 50 | - name: Remove CD container 51 | if: always() 52 | working-directory: ${{github.workspace}} 53 | run: | 54 | bash scripts/remove_container.sh ${JOBNAME} 55 | -------------------------------------------------------------------------------- /.github/workflows/scripts/build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | GLT_ROOT_DIR=$(dirname $(dirname $(dirname $(dirname "$(realpath "$0")")))) 6 | GITHUB_REF=$1 7 | 8 | rm -rf /usr/local/cuda 9 | ln -s /usr/local/cuda-12.1 /usr/local/cuda 10 | 11 | cd $GLT_ROOT_DIR 12 | 13 | if [[ "$GITHUB_REF" =~ ^"refs/tags/" ]]; then 14 | export GITHUB_TAG_REF="$GITHUB_REF" 15 | fi 16 | 17 | if [ -z "$GITHUB_TAG_REF" ]; then 18 | echo "Not on a tag, won't deploy to pypi" 19 | else 20 | PYABIS="cp38-cp38 cp39-cp39 cp310-cp310 cp311-cp311" 21 | for abi in $PYABIS; do 22 | PYABI=$abi bash .github/workflows/scripts/build_wheel.sh 23 | done 24 | fi 25 | -------------------------------------------------------------------------------- /.github/workflows/scripts/build_wheel.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | GLT_ROOT_DIR=$(dirname $(dirname $(dirname $(dirname "$(realpath "$0")")))) 6 | PYBIN=/opt/python/${PYABI}/bin 7 | 8 | cd $GLT_ROOT_DIR 9 | 10 | set -x 11 | RELEASE=TRUE WITH_CUDA=ON ${PYBIN}/python setup.py bdist_wheel 12 | 13 | # Bundle external shared libraries into the wheels 14 | for whl in dist/*-linux*.whl; do 15 | ${PYBIN}/auditwheel repair "$whl" -w dist/ --plat manylinux2014_x86_64 --exclude libtorch_cpu.so --exclude libc10.so --exclude libtorch_python.so --exclude libtorch.so 16 | done 17 | 18 | rm dist/*-linux*.whl 19 | -------------------------------------------------------------------------------- /.github/workflows/scripts/upload_pypi.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | GLT_ROOT_DIR=$(dirname $(dirname $(dirname $(dirname "$(realpath "$0")")))) 6 | cd $GLT_ROOT_DIR 7 | 8 | echo "********" 9 | echo "Upload packages:" 10 | ls ${GLT_ROOT_DIR}/dist/ 11 | echo "********" 12 | 13 | echo "[distutils]" > ~/.pypirc 14 | echo "index-servers =" >> ~/.pypirc 15 | echo " pypi" >> ~/.pypirc 16 | echo "[pypi]" >> ~/.pypirc 17 | echo "repository=https://upload.pypi.org/legacy/" >> ~/.pypirc 18 | echo "username=__token__" >> ~/.pypirc 19 | echo "password=$PYPI_PWD" >> ~/.pypirc 20 | 21 | pip3 install twine 22 | python3 -m twine upload -r pypi --skip-existing ${GLT_ROOT_DIR}/dist/* 23 | rm -rf ~/.pypirc 24 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | built/ 2 | build/ 3 | _build/ 4 | /dist/ 5 | /graphlearn_torch.egg-info/ 6 | /data/ 7 | /examples/data/ 8 | CMakeFiles/ 9 | *.cmake 10 | __pycache__/ 11 | .vscode/ 12 | CMakeCache.txt 13 | Makefile 14 | detect_cuda_* 15 | .idea/ 16 | cmake-build-release/ 17 | cmake-build-debug/ 18 | /graphlearn_torch/python/*.so -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "third_party/googletest/googletest"] 2 | path = third_party/googletest/googletest 3 | url = https://github.com/google/googletest.git 4 | -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | # .readthedocs.yaml 2 | # Read the Docs configuration file 3 | # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details 4 | 5 | # Required 6 | version: 2 7 | 8 | build: 9 | os: ubuntu-20.04 10 | tools: 11 | python: "3.8" 12 | 13 | sphinx: 14 | configuration: docs/conf.py 15 | 16 | python: 17 | install: 18 | - requirements: docs/requirements.txt -------------------------------------------------------------------------------- /benchmarks/api/bench_dist_config.yml: -------------------------------------------------------------------------------- 1 | nodes: 2 | - 0.0.0.0 3 | - 1.1.1.1 4 | ports: [22,22] 5 | python_bins: 6 | - /path/to/python 7 | - /path/to/python # path to python with GLT envs 8 | usernames: 9 | - root 10 | - root 11 | dataset: products 12 | node_ranks: [0, 1] 13 | dst_paths: 14 | - /path/to/gl_torch 15 | - /path/to/gl_torch 16 | visible_devices: 17 | - 0,1,2,3 18 | - 4,5,6,7 19 | -------------------------------------------------------------------------------- /dockerfiles/graphlearn-torch-dev.Dockerfile: -------------------------------------------------------------------------------- 1 | ARG GLT_CUDA_VERSION=12.1 2 | ARG GLT_TORCH_VERSION=2.3.0 3 | 4 | FROM nvidia/cuda:${GLT_CUDA_VERSION}.0-devel-ubuntu22.04 5 | 6 | RUN apt update && apt install python3-pip vim git cmake -y 7 | RUN ln -s python3 /usr/bin/python 8 | 9 | ARG GLT_TORCH_VERSION GLT_CUDA_VERSION 10 | 11 | RUN CUDA_WHEEL_VERSION=$(echo "$GLT_CUDA_VERSION" | sed 's/\.//g') \ 12 | && pip install torch==${GLT_TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${CUDA_WHEEL_VERSION} \ 13 | && pip install torch_geometric \ 14 | && pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv \ 15 | -f https://data.pyg.org/whl/torch-${GLT_TORCH_VERSION}+cu${CUDA_WHEEL_VERSION}.html 16 | 17 | -------------------------------------------------------------------------------- /dockerfiles/graphlearn-torch-wheel.Dockerfile: -------------------------------------------------------------------------------- 1 | ARG GLT_TORCH_VERSION=2.3.0 2 | ARG CUDA_WHEEL_VERSION=121 3 | 4 | FROM pytorch/manylinux-cuda$CUDA_WHEEL_VERSION 5 | 6 | ARG GLT_TORCH_VERSION CUDA_WHEEL_VERSION 7 | ENV ABIS="cp38-cp38 cp39-cp39 cp310-cp310 cp311-cp311" 8 | 9 | RUN set -ex; \ 10 | for abi in $ABIS; do \ 11 | PYBIN=/opt/python/${abi}/bin; \ 12 | ${PYBIN}/pip install ninja scipy; \ 13 | ${PYBIN}/pip install torch==${GLT_TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${CUDA_WHEEL_VERSION}; \ 14 | ${PYBIN}/pip install torch_geometric; \ 15 | ${PYBIN}/pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv \ 16 | -f https://data.pyg.org/whl/torch-${GLT_TORCH_VERSION}+cu${CUDA_WHEEL_VERSION}.html; \ 17 | ${PYBIN}/pip install auditwheel; \ 18 | done -------------------------------------------------------------------------------- /dockerfiles/gs-graphlearn-torch-dev.Dockerfile: -------------------------------------------------------------------------------- 1 | ARG REGISTRY=registry.cn-hongkong.aliyuncs.com 2 | ARG GLT_TORCH_VERSION=2.3.0 3 | ARG GS_VERSION=latest 4 | 5 | FROM $REGISTRY/graphscope/graphscope-dev:$GS_VERSION 6 | 7 | ARG GLT_TORCH_VERSION 8 | 9 | RUN pip install torch==${GLT_TORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu \ 10 | && pip install torch_geometric \ 11 | && pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv \ 12 | -f https://data.pyg.org/whl/torch-${GLT_TORCH_VERSION}+cpu.html 13 | 14 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = sphinx-build 7 | PAPER = 8 | BUILDDIR = _build 9 | 10 | # User-friendly check for sphinx-build 11 | ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1) 12 | $(error The '$(SPHINXBUILD)' command was not found. Make sure you have Sphinx installed, then set the SPHINXBUILD environment variable to point to the full path of the '$(SPHINXBUILD)' executable. Alternatively you can add the directory with the executable to your PATH. If you don't have Sphinx installed, grab it from http://sphinx-doc.org/) 13 | endif 14 | 15 | # Internal variables. 16 | PAPEROPT_a4 = -D latex_paper_size=a4 17 | PAPEROPT_letter = -D latex_paper_size=letter 18 | ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . 19 | # the i18n builder cannot share the environment and doctrees with the others 20 | I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . 21 | 22 | .PHONY: help 23 | help: 24 | @echo "Please use \`make ' where is one of" 25 | @echo " html to make standalone HTML files" 26 | @echo " dirhtml to make HTML files named index.html in directories" 27 | @echo " singlehtml to make a single large HTML file" 28 | @echo " pickle to make pickle files" 29 | @echo " json to make JSON files" 30 | @echo " htmlhelp to make HTML files and a HTML help project" 31 | @echo " qthelp to make HTML files and a qthelp project" 32 | @echo " applehelp to make an Apple Help Book" 33 | @echo " devhelp to make HTML files and a Devhelp project" 34 | @echo " epub to make an epub" 35 | @echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter" 36 | @echo " latexpdf to make LaTeX files and run them through pdflatex" 37 | @echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx" 38 | @echo " text to make text files" 39 | @echo " man to make manual pages" 40 | @echo " texinfo to make Texinfo files" 41 | @echo " info to make Texinfo files and run them through makeinfo" 42 | @echo " gettext to make PO message catalogs" 43 | @echo " changes to make an overview of all changed/added/deprecated items" 44 | @echo " xml to make Docutils-native XML files" 45 | @echo " pseudoxml to make pseudoxml-XML files for display purposes" 46 | @echo " linkcheck to check all external links for integrity" 47 | @echo " doctest to run all doctests embedded in the documentation (if enabled)" 48 | @echo " coverage to run coverage check of the documentation (if enabled)" 49 | 50 | .PHONY: clean 51 | clean: 52 | rm -rf $(BUILDDIR)/* 53 | 54 | .PHONY: html 55 | html: 56 | $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html 57 | @echo 58 | @echo "Build finished. The HTML pages are in $(BUILDDIR)/html." 59 | -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- 1 | # Documents 2 | ## Generate documents 3 | 4 | 1. install requirements 5 | 6 | ``` 7 | cd docs 8 | pip install -r requirements.txt 9 | ``` 10 | 11 | 2. generate htmls 12 | 13 | ``` 14 | # make clean 15 | make html 16 | ``` 17 | 18 | ## Update API Reference 19 | 20 | If a new module added or an existed module deleted, then update 21 | the .rst files in docs/apis 22 | 23 | ``` 24 | sphinx-apidoc -o ./apis ../graphlearn_torch/python/{module_name} 25 | # update docs/apis/modules 26 | make clean 27 | make html 28 | ``` -------------------------------------------------------------------------------- /docs/apis/channel.rst: -------------------------------------------------------------------------------- 1 | graphlearn_torch.channel 2 | ======================== 3 | 4 | base 5 | ---- 6 | 7 | .. automodule:: graphlearn_torch.channel.base 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | mp\_channel 13 | ----------- 14 | 15 | .. automodule:: graphlearn_torch.channel.mp_channel 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | 20 | remote\_channel 21 | --------------- 22 | 23 | .. automodule:: graphlearn_torch.channel.remote_channel 24 | :members: 25 | :undoc-members: 26 | :show-inheritance: 27 | 28 | shm\_channel 29 | ------------ 30 | 31 | .. automodule:: graphlearn_torch.channel.shm_channel 32 | :members: 33 | :undoc-members: 34 | :show-inheritance: 35 | -------------------------------------------------------------------------------- /docs/apis/data.rst: -------------------------------------------------------------------------------- 1 | graphlearn_torch.data 2 | ====================== 3 | 4 | dataset 5 | ------- 6 | 7 | .. automodule:: graphlearn_torch.data.dataset 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | feature 13 | ------- 14 | 15 | .. automodule:: graphlearn_torch.data.feature 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | 20 | graph 21 | ----- 22 | 23 | .. automodule:: graphlearn_torch.data.graph 24 | :members: 25 | :undoc-members: 26 | :show-inheritance: 27 | 28 | reorder 29 | ------- 30 | 31 | .. automodule:: graphlearn_torch.data.reorder 32 | :members: 33 | :undoc-members: 34 | :show-inheritance: 35 | 36 | 37 | table\_dataset 38 | -------------- 39 | 40 | .. automodule:: graphlearn_torch.data.table_dataset 41 | :members: 42 | :undoc-members: 43 | :show-inheritance: 44 | 45 | unified\_tensor 46 | --------------- 47 | 48 | .. automodule:: graphlearn_torch.data.unified_tensor 49 | :members: 50 | :undoc-members: 51 | :show-inheritance: 52 | 53 | vineyard\_utils 54 | --------------- 55 | 56 | .. automodule:: graphlearn_torch.data.vineyard_utils 57 | :members: 58 | :undoc-members: 59 | :show-inheritance: 60 | -------------------------------------------------------------------------------- /docs/apis/distributed.rst: -------------------------------------------------------------------------------- 1 | graphlearn_torch.distributed 2 | =============================== 3 | 4 | dist\_client 5 | ------------ 6 | 7 | .. automodule:: graphlearn_torch.distributed.dist_client 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | dist\_context 13 | ------------- 14 | 15 | .. automodule:: graphlearn_torch.distributed.dist_context 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | 20 | dist\_dataset 21 | ------------- 22 | 23 | .. automodule:: graphlearn_torch.distributed.dist_dataset 24 | :members: 25 | :undoc-members: 26 | :show-inheritance: 27 | 28 | dist\_feature 29 | ------------- 30 | 31 | .. automodule:: graphlearn_torch.distributed.dist_feature 32 | :members: 33 | :undoc-members: 34 | :show-inheritance: 35 | 36 | dist\_graph 37 | ----------- 38 | 39 | .. automodule:: graphlearn_torch.distributed.dist_graph 40 | :members: 41 | :undoc-members: 42 | :show-inheritance: 43 | 44 | dist\_neighbor\_loader 45 | ---------------------- 46 | 47 | .. automodule:: graphlearn_torch.distributed.dist_neighbor_loader 48 | :members: 49 | :undoc-members: 50 | :show-inheritance: 51 | 52 | 53 | dist\_neighbor\_sampler 54 | ----------------------- 55 | 56 | .. automodule:: graphlearn_torch.distributed.dist_neighbor_sampler 57 | :members: 58 | :undoc-members: 59 | :show-inheritance: 60 | 61 | dist\_options 62 | ------------- 63 | 64 | .. automodule:: graphlearn_torch.distributed.dist_options 65 | :members: 66 | :undoc-members: 67 | :show-inheritance: 68 | 69 | dist\_sampling\_producer 70 | ------------------------ 71 | 72 | .. automodule:: graphlearn_torch.distributed.dist_sampling_producer 73 | :members: 74 | :undoc-members: 75 | :show-inheritance: 76 | 77 | dist\_server 78 | ------------ 79 | 80 | .. automodule:: graphlearn_torch.distributed.dist_server 81 | :members: 82 | :undoc-members: 83 | :show-inheritance: 84 | 85 | event\_loop 86 | ----------- 87 | 88 | .. automodule:: graphlearn_torch.distributed.event_loop 89 | :members: 90 | :undoc-members: 91 | :show-inheritance: 92 | 93 | rpc 94 | --- 95 | 96 | .. automodule:: graphlearn_torch.distributed.rpc 97 | :members: 98 | :undoc-members: 99 | :show-inheritance: 100 | -------------------------------------------------------------------------------- /docs/apis/loader.rst: -------------------------------------------------------------------------------- 1 | graphlearn_torch.loader 2 | ============================== 3 | 4 | neighbor\_loader 5 | ---------------- 6 | 7 | .. automodule:: graphlearn_torch.loader.neighbor_loader 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | transform 13 | --------- 14 | 15 | .. automodule:: graphlearn_torch.loader.transform 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | -------------------------------------------------------------------------------- /docs/apis/modules.rst: -------------------------------------------------------------------------------- 1 | .. toctree:: 2 | :maxdepth: 4 3 | 4 | channel 5 | data 6 | distributed 7 | loader 8 | partition 9 | sampler 10 | utils 11 | -------------------------------------------------------------------------------- /docs/apis/partition.rst: -------------------------------------------------------------------------------- 1 | graphlearn_torch.partition 2 | ========================== 3 | 4 | feature\_partition 5 | ------------------- 6 | 7 | .. automodule:: graphlearn_torch.partition.feature_partition 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | graph\_partition 13 | ---------------- 14 | 15 | .. automodule:: graphlearn_torch.partition.graph_partition 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | -------------------------------------------------------------------------------- /docs/apis/sampler.rst: -------------------------------------------------------------------------------- 1 | graphlearn_torch.sampler 2 | ======================== 3 | 4 | base 5 | ---- 6 | 7 | .. automodule:: graphlearn_torch.sampler.base 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | negative\_sampler 13 | ------------------ 14 | 15 | .. automodule:: graphlearn_torch.sampler.negative_sampler 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | 20 | neighbor\_sampler 21 | ------------------ 22 | 23 | .. automodule:: graphlearn_torch.sampler.neighbor_sampler 24 | :members: 25 | :undoc-members: 26 | :show-inheritance: 27 | -------------------------------------------------------------------------------- /docs/apis/utils.rst: -------------------------------------------------------------------------------- 1 | graphlearn_torch.utils 2 | ====================== 3 | 4 | common 5 | ------- 6 | 7 | .. automodule:: graphlearn_torch.utils.common 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | device 13 | ------- 14 | 15 | .. automodule:: graphlearn_torch.utils.device 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | 20 | exit\_status 21 | ------------- 22 | 23 | .. automodule:: graphlearn_torch.utils.exit_status 24 | :members: 25 | :undoc-members: 26 | :show-inheritance: 27 | 28 | mixin 29 | ------ 30 | 31 | .. automodule:: graphlearn_torch.utils.mixin 32 | :members: 33 | :undoc-members: 34 | :show-inheritance: 35 | 36 | singleton 37 | --------- 38 | 39 | .. automodule:: graphlearn_torch.utils.singleton 40 | :members: 41 | :undoc-members: 42 | :show-inheritance: 43 | 44 | tensor 45 | ------ 46 | 47 | .. automodule:: graphlearn_torch.utils.tensor 48 | :members: 49 | :undoc-members: 50 | :show-inheritance: 51 | 52 | units 53 | ------ 54 | 55 | .. automodule:: graphlearn_torch.utils.units 56 | :members: 57 | :undoc-members: 58 | :show-inheritance: 59 | -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # For the full list of built-in configuration values, see the documentation: 4 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 5 | 6 | # -- Project information ----------------------------------------------------- 7 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information 8 | 9 | import datetime 10 | 11 | import graphlearn_torch 12 | 13 | project = 'graphlearn-for-pytorch' 14 | author = 'Alibaba-inc' 15 | copyright = f'{datetime.datetime.now().year}, {author}' 16 | 17 | # -- General configuration --------------------------------------------------- 18 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration 19 | 20 | from recommonmark.parser import CommonMarkParser 21 | source_parsers = { 22 | '.md': CommonMarkParser, 23 | } 24 | source_suffix = ['.rst', '.md'] 25 | 26 | extensions = [ 27 | 'sphinx.ext.autodoc', 28 | 'sphinx.ext.autosummary', 29 | 'sphinx.ext.intersphinx', 30 | 'sphinx.ext.mathjax', 31 | 'sphinx.ext.napoleon', 32 | 'sphinx.ext.viewcode', 33 | "myst_parser", 34 | ] 35 | 36 | add_module_names = False 37 | autodoc_member_order = 'bysource' 38 | 39 | templates_path = ['_templates'] 40 | exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] 41 | 42 | # -- Options for HTML output ------------------------------------------------- 43 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output 44 | 45 | html_theme = 'sphinx_rtd_theme' 46 | html_static_path = ['_static'] 47 | 48 | pygments_style = "sphinx" 49 | 50 | # If true, `todo` and `todoList` produce output, else they produce nothing. 51 | todo_include_todos = False -------------------------------------------------------------------------------- /docs/contribute.md: -------------------------------------------------------------------------------- 1 | # Contribute to GLT -------------------------------------------------------------------------------- /docs/faq.md: -------------------------------------------------------------------------------- 1 | # Frequently Asked Questions(FAQ) 2 | 3 | 1. Bus Error 4 | 5 | GLT uses shared memory for sharing data between processes. If there is a bus error, 6 | make sure your memory can hold the data and you need to set the shared memory to a 7 | larger size, e.g 128G: 8 | ```shell 9 | echo 128000000000 > /proc/sys/kernel/shmmax 10 | mount -o remount,size=128G /dev/shm 11 | ``` -------------------------------------------------------------------------------- /docs/figures/arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/arch.png -------------------------------------------------------------------------------- /docs/figures/dist_arch_server_client_mode.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/dist_arch_server_client_mode.png -------------------------------------------------------------------------------- /docs/figures/dist_arch_worker_mode.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/dist_arch_worker_mode.png -------------------------------------------------------------------------------- /docs/figures/glt_dgl.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/glt_dgl.png -------------------------------------------------------------------------------- /docs/figures/replica_per_dg.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/replica_per_dg.png -------------------------------------------------------------------------------- /docs/figures/replica_per_gpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/replica_per_gpu.png -------------------------------------------------------------------------------- /docs/figures/scale_out.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/scale_out.png -------------------------------------------------------------------------------- /docs/figures/scale_up.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/scale_up.png -------------------------------------------------------------------------------- /docs/figures/uni_tensor.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/docs/figures/uni_tensor.png -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- 1 | .. graphlearn_torch documentation master file, created by 2 | sphinx-quickstart on Mon Jan 30 07:33:10 2023. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | Welcome to GraphLearn-for-PyTorch's documentation! 7 | ================================================== 8 | **GraphLearn-for-PyTorch(GLT)** is a graph learning library for PyTorch that 9 | makes distributed GNN training and inference easy and efficient. 10 | It leverages the power of GPUs to accelerate graph sampling and utilizes 11 | UVA to reduce the conversion and copying of features of vertices and edges. 12 | For large-scale graphs, it supports distributed training on multiple GPUs or 13 | multiple machines through fast distributed sampling and feature lookup. 14 | Additionally, it provides flexible deployment for distributed training to meet 15 | different requirements. 16 | 17 | .. toctree:: 18 | :maxdepth: 2 19 | :caption: Installation 20 | 21 | install/install 22 | 23 | .. toctree:: 24 | :maxdepth: 2 25 | :caption: Get Started 26 | 27 | get_started/graph_basic 28 | get_started/node_class 29 | get_started/link_pred 30 | get_started/hetero 31 | get_started/dist_train 32 | 33 | .. toctree:: 34 | :maxdepth: 2 35 | :caption: Tutorials 36 | 37 | tutorial/basic_object 38 | tutorial/graph_ops 39 | tutorial/dist 40 | tutorial/aliyun 41 | 42 | 43 | .. toctree:: 44 | :maxdepth: 3 45 | :caption: API Reference 46 | 47 | apis/modules 48 | 49 | .. toctree:: 50 | :maxdepth: 2 51 | :caption: Contribute to GLT 52 | 53 | contribute 54 | 55 | .. toctree:: 56 | :maxdepth: 2 57 | :caption: FAQ 58 | 59 | faq 60 | 61 | Indices and tables 62 | ================== 63 | 64 | * :ref:`genindex` 65 | * :ref:`modindex` 66 | * :ref:`search` 67 | -------------------------------------------------------------------------------- /docs/install/install.md: -------------------------------------------------------------------------------- 1 | # Installation 2 | 3 | ### Requirements 4 | - cuda 5 | - python>=3.6 6 | - torch(PyTorch) 7 | - torch_geometric, torch_scatter, torch_sparse. Please refer to [PyG](https://github.com/pyg-team/pytorch_geometric) for installation. 8 | ### Pip Wheels 9 | 10 | ``` 11 | # glibc>=2.14, torch>=1.13 12 | pip install graphlearn-torch 13 | ``` 14 | 15 | ### Build from source 16 | 17 | #### Install Dependencies 18 | ```shell 19 | git submodule update --init 20 | ./install_dependencies.sh 21 | ``` 22 | 23 | #### Python 24 | 1. Build 25 | ``` shell 26 | python setup.py bdist_wheel 27 | pip install dist/* 28 | ``` 29 | 2. UT 30 | ``` shell 31 | ./scripts/run_python_ut.sh 32 | ``` 33 | 34 | #### C++ 35 | If you need to test C++ operations, you can only build the C++ part. 36 | 37 | 1. Build 38 | ``` shell 39 | cmake . 40 | make -j 41 | ``` 42 | 2. UT 43 | ``` shell 44 | ./scripts/run_cpp_ut.sh 45 | ``` 46 | -------------------------------------------------------------------------------- /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=. 11 | set BUILDDIR=_build 12 | 13 | %SPHINXBUILD% >NUL 2>NUL 14 | if errorlevel 9009 ( 15 | echo. 16 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 17 | echo.installed, then set the SPHINXBUILD environment variable to point 18 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 19 | echo.may add the Sphinx directory to PATH. 20 | echo. 21 | echo.If you don't have Sphinx installed, grab it from 22 | echo.https://www.sphinx-doc.org/ 23 | exit /b 1 24 | ) 25 | 26 | if "%1" == "" goto help 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/requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | torch_geometric 3 | --extra-index-url https://download.pytorch.org/whl/cpu 4 | torch==1.12.0+cpu 5 | --extra-index-url https://download.pytorch.org/whl/cpu 6 | torchvision==0.13.0+cpu 7 | --extra-index-url https://download.pytorch.org/whl/cpu 8 | torchaudio==0.12.0 9 | https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_sparse-0.6.16%2Bpt112cpu-cp38-cp38-linux_x86_64.whl 10 | https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_scatter-2.1.0%2Bpt112cpu-cp38-cp38-linux_x86_64.whl 11 | https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_cluster-1.6.0%2Bpt112cpu-cp38-cp38-linux_x86_64.whl 12 | sphinx 13 | https://graphlearn.oss-cn-hangzhou.aliyuncs.com/cpu/graphlearn_torch-0.2.0a1-cp38-cp38-linux_x86_64.whl 14 | sphinx_rtd_theme 15 | recommonmark 16 | markdown 17 | sphinx-markdown-tables 18 | myst-parser 19 | -------------------------------------------------------------------------------- /examples/README.md: -------------------------------------------------------------------------------- 1 | # GNN examples 2 | 3 | 1. Multi-processing Feature usage. 4 | 5 | ``` 6 | python feature_mp.py 7 | ``` 8 | 9 | 2. Single GPU basic GraphSAGE on OGBN-Products similar to PyG. 10 | This example is implemeted using `NeighborSampler`. 11 | 12 | ``` 13 | python train_sage_ogbn_products.py 14 | ``` 15 | 16 | 3. Multiple GPUs GraphSAGE on OGBN-papers100M. 17 | ``` 18 | python multi_gpu/train_sage_ogbn_papers100m.py 19 | ``` 20 | 21 | 4. Heterogeneous GraphSage examples on OGBN-MAG compatible with PyG. 22 | 23 | ``` 24 | # single GPU 25 | python hetero/train_hgt_mag.py 26 | # multi-GPUs 27 | python hetero/train_hgt_mag_mp.py 28 | ``` 29 | 30 | 5. Training on PAI. 31 | 32 | see pai/README.md 33 | 34 | 6. Distributed (multi nodes) examples. 35 | see distributed/README.md -------------------------------------------------------------------------------- /examples/distributed/dist_train_sage_sup_config.yml: -------------------------------------------------------------------------------- 1 | nodes: 2 | - 0.0.0.0 3 | - 1.1.1.1 4 | ports: [22,22] 5 | python_bins: 6 | - /path/to/python 7 | - /path/to/python # path to python with GLT's env 8 | usernames: 9 | - root 10 | - root 11 | dataset: ogbn-products 12 | in_channel: 100 13 | out_channel: 47 14 | node_ranks: [0, 1] 15 | dst_paths: 16 | - /path/to/gl_torch 17 | - /path/to/gl_torch 18 | visible_devices: 19 | - 0,1,2,3 20 | - 4,5,6,7 21 | -------------------------------------------------------------------------------- /examples/distributed/run_dist_train_sage_sup.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import yaml 17 | import argparse 18 | 19 | import paramiko 20 | import click 21 | 22 | if __name__ == "__main__": 23 | parser = argparse.ArgumentParser('Run DistRandomSampler benchmarks.') 24 | parser.add_argument('--config', type=str, default='dist_train_sage_sup_config.yml', 25 | help='paths to configuration file for benchmarks') 26 | parser.add_argument('--epochs', type=int, default=1, 27 | help='repeat epochs for sampling') 28 | parser.add_argument('--batch_size', type=int, default=2048, 29 | help='batch size for sampling') 30 | parser.add_argument('--master_addr', type=str, default='0.0.0.0', 31 | help='master ip address for synchronization across all training nodes') 32 | parser.add_argument('--master_port', type=str, default='11345', 33 | help='port for synchronization across all training nodes') 34 | args = parser.parse_args() 35 | 36 | config = open(args.config, 'r') 37 | config = yaml.safe_load(config) 38 | dataset = config['dataset'] 39 | ip_list, port_list, username_list = config['nodes'], config['ports'], config['usernames'] 40 | dst_path_list = config['dst_paths'] 41 | node_ranks = config['node_ranks'] 42 | num_nodes = len(node_ranks) 43 | visible_devices = config['visible_devices'] 44 | python_bins = config['python_bins'] 45 | num_cores = len(visible_devices[0].split(',')) 46 | in_channel = str(config['in_channel']) 47 | out_channel = str(config['out_channel']) 48 | 49 | dataset_path = "../../data/" 50 | passwd_dict = {} 51 | for username, ip in zip(username_list, ip_list): 52 | passwd_dict[ip+username] = click.prompt('passwd for '+username+'@'+ip, 53 | hide_input=True) 54 | for username, ip, port, dst, noderk, device, pythonbin in zip( 55 | username_list, 56 | ip_list, 57 | port_list, 58 | dst_path_list, 59 | node_ranks, 60 | visible_devices, 61 | python_bins, 62 | ): 63 | trans = paramiko.Transport((ip, port)) 64 | trans.connect(username=username, password=passwd_dict[ip+username]) 65 | ssh = paramiko.SSHClient() 66 | ssh._transport = trans 67 | 68 | to_dist_dir = 'cd '+dst+'/examples/distributed/ ' 69 | exec_example = "tmux new -d 'CUDA_VISIBLE_DEVICES="+device+" "+pythonbin+" dist_train_sage_supervised.py --dataset="+dataset+" --dataset_root_dir=../../data/"+dataset+" --in_channel="+in_channel+" --out_channel="+out_channel+" --node_rank="+str(noderk)+" --num_dataset_partitions="+str(num_nodes)+" --num_nodes="+str(num_nodes)+" --num_training_procs="+str(num_cores)+" --master_addr="+args.master_addr+" --training_pg_master_port="+args.master_port+" --train_loader_master_port="+str(int(args.master_port)+1)+" --test_loader_master_port="+str(int(args.master_port)+2)+" --batch_size="+str(args.batch_size)+" --epochs="+str(args.epochs) 70 | 71 | print(to_dist_dir + ' && '+ exec_example + " '") 72 | stdin, stdout, stderr = ssh.exec_command(to_dist_dir+' && '+exec_example+" '", bufsize=1) 73 | print(stdout.read().decode()) 74 | print(stderr.read().decode()) 75 | ssh.close() 76 | -------------------------------------------------------------------------------- /examples/feature_mp.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import torch 17 | import graphlearn_torch as glt 18 | 19 | import torch.multiprocessing as mp 20 | 21 | def test_feature(rank, world_size, feature): 22 | torch.cuda.set_device(rank) 23 | assert list(feature.shape) == [128*3, 128] 24 | input = torch.tensor([10, 20, 200, 210, 300, 310], dtype=torch.int64, 25 | device= torch.device(rank)) 26 | attr = torch.ones(2, 128, dtype=torch.float32, 27 | device=torch.device(rank)) 28 | res = torch.cat((attr, attr*2, attr*3), 0) 29 | assert glt.utils.tensor_equal_with_device(feature[input], res) 30 | 31 | if __name__ == "__main__": 32 | print('Use GPU 0 and GPU 1 for multiprocessing Feature test.') 33 | world_size = 2 34 | attr = torch.ones(128, 128, dtype=torch.float32) 35 | tensor = torch.cat([attr, attr*2, attr*3], 0) 36 | 37 | rows = torch.cat([torch.arange(128*3), 38 | torch.randint(128, (128*3,)), 39 | torch.randint(128*2, (128*3,))]) 40 | cols = torch.cat([torch.randint(128*3, (128*3,)), 41 | torch.randint(128*3, (128*3,)), 42 | torch.randint(128*3, (128*3,))]) 43 | csr_topo = glt.data.Topology(edge_index=torch.stack([rows, cols])) 44 | 45 | device_group_list = [glt.data.DeviceGroup(0, [0]), 46 | glt.data.DeviceGroup(1, [1])] 47 | #device_group_list = [glt.data.DeviceGroup(0, [0, 1])] 48 | split_ratio = 0.8 # [0, 1] 49 | cpu_tensor, id2index = glt.data.sort_by_in_degree(tensor, split_ratio, csr_topo) 50 | feature = glt.data.Feature(cpu_tensor, id2index, split_ratio, device_group_list, 0) 51 | mp.spawn(test_feature, 52 | args=(world_size, feature), 53 | nprocs=world_size, 54 | join=True) -------------------------------------------------------------------------------- /examples/gpt/README.md: -------------------------------------------------------------------------------- 1 | # Using GPT to Reason on Graphs 2 | 3 | This simple example shows how to leverage GPT to make inference on large graphs. 4 | 5 | ### 1. Prepare the graph dataset & environment 6 | In this example, we use the dataset 7 | [arxiv_2023](https://github.com/TRAIS-Lab/LLM-Structured-Data/tree/main/dataset/arxiv_2023) 8 | and download it to the path `../data/arxiv_2023`. 9 | 10 | Then, export your OPENAI_API_KEY as the environment variable in your shell: 11 | 12 | ```bash 13 | export OPENAI_API_KEY='YOUR_API_KEY' 14 | ``` 15 | 16 | ### 2. Run the code 17 | Configure the parameters in the data loader and run the code. 18 | ```bash 19 | python arxiv.py 20 | ``` 21 | This example tests the inference performance of GPT on a large graph with the link prediction task as the default task. 22 | 23 | First, we sample a 2-hop ego-subgraph from the original graph. The subgraph is sampled by the `LinkNeighborLoader` with a mini-batch sampler that samples a fixed number of neighbors for each edge and is formed as PyG's `edge_index`. 24 | 25 | Then, the sampled subgraph along with node features (e.g. in this example the title for the paper node) is fed into GPT to infer whether a requested edge is in the original graph or not. 26 | 27 | **Note**: GPT has a limitation on its context size, and thus limits the size of the sampled subgraph, which is determined by the parameter `num_neighbors` in the data loader. If the sampled subgraph is too large, please try to decrease the `num_neighbors` to reduce the size of the subgraph. 28 | 29 | ### Appendix: 30 | 1. **Dataset**: You can also use other datasets and modify the preprocessing code, but don't forget to transform the graph format into PyG's `edge_index`. 31 | 2. **Prompts**: Use the parameter `reason: Bool` to decide whether to see the reasoning process of GPT. You can also design your own prompts to make inference on graphs instead of using our template. 32 | 3. **Node classification**: We also provide a template prompt for node classification task, and design your method to leverage the label informatiion. 33 | **Note**: We've tried directly passing the labels for nodes in an ego-subgraph to predict the label of the center node, and GPT's prediction behavior in this case is close to voting via neighbors. 34 | -------------------------------------------------------------------------------- /examples/gpt/arxiv.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import time 17 | import torch 18 | 19 | from tqdm import tqdm 20 | 21 | import graphlearn_torch as glt 22 | from utils import get_gpt_response, link_prediction 23 | 24 | 25 | def run(glt_ds, raw_text, reason): 26 | neg_sampling = glt.sampler.NegativeSampling('binary') 27 | train_loader = glt.loader.LinkNeighborLoader(glt_ds, 28 | [12, 6], 29 | neg_sampling=neg_sampling, 30 | batch_size=2, 31 | drop_last=True, 32 | shuffle=True, 33 | device=torch.device('cpu')) 34 | print(f'Building graphlearn_torch NeighborLoader Done.') 35 | 36 | for batch in tqdm(train_loader): 37 | batch_titles = raw_text[batch.node] 38 | if batch.edge_index.shape[1] < 5: 39 | continue 40 | 41 | # print(batch) 42 | # print(batch.edge_label_index) 43 | message = link_prediction(batch, batch_titles, reason=reason) 44 | 45 | # print(message) 46 | response = get_gpt_response( 47 | message=message 48 | ) 49 | 50 | print(f"response: {response}") 51 | 52 | 53 | if __name__ == '__main__': 54 | import pandas as pd 55 | root = '../data/arxiv_2023/raw/' 56 | titles = pd.read_csv(root + "titles.csv.gz").to_numpy() 57 | ids = torch.from_numpy(pd.read_csv(root + "ids.csv.gz").to_numpy()) 58 | edge_index = torch.from_numpy(pd.read_csv(root + "edges.csv.gz").to_numpy()) 59 | 60 | print('Build graphlearn_torch dataset...') 61 | start = time.time() 62 | glt_dataset = glt.data.Dataset() 63 | glt_dataset.init_graph( 64 | edge_index=edge_index.T, 65 | graph_mode='CPU', 66 | directed=True 67 | ) 68 | glt_dataset.init_node_features( 69 | node_feature_data=ids, 70 | sort_func=glt.data.sort_by_in_degree, 71 | split_ratio=0 72 | ) 73 | 74 | print(f'Build graphlearn_torch csr_topo and feature cost {time.time() - start} s.') 75 | 76 | run(glt_dataset, titles, reason=False) 77 | -------------------------------------------------------------------------------- /examples/gpt/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from openai import OpenAI 17 | 18 | 19 | def get_gpt_response(message, model="gpt-4-1106-preview"): 20 | client = OpenAI() 21 | chat_completion = client.chat.completions.create( 22 | messages=[ 23 | { 24 | "role" : "user", 25 | "content": message, 26 | } 27 | ], 28 | model=model, 29 | ) 30 | return chat_completion.choices[0].message.content 31 | 32 | 33 | def node_classification(batch): 34 | message = "This is a directed subgraph of arxiv citation network with " + str(batch.x.shape[0]) + " nodes numbered from 0 to " + str(batch.x.shape[0]-1) + ".\n" 35 | message += "The subgraph has " + str(batch.edge_index.shape[1]) + " edges.\n" 36 | for i in range(1, batch.x.shape[0]): 37 | feature_str = ','.join(f'{it:.3f}' for it in batch.x[i].tolist()) 38 | message += "The feature of node " + str(i) + " is [" + feature_str + "] " 39 | message += "and the node label is " + str(batch.y[i].item()) + ".\n" 40 | message += "The edges of the subgraph are " + str(batch.edge_index.T.tolist()) + ' where the first number indicates source node and the second destination node.\n' 41 | message += "Question: predict the label for node 0, whose feature is [" + ','.join(f'{it:.3f}' for it in batch.x[0].tolist()) + "]. Give the label only and don't show any reasoning process.\n\n" 42 | 43 | return message 44 | 45 | 46 | def link_prediction(batch, titles, reason=False): 47 | message = "This is a directed subgraph of arxiv citation network with " + str(batch.x.shape[0]) + " nodes numbered from 0 to " + str(batch.x.shape[0]-1) + ".\n" 48 | graph = batch.edge_index.T.unique(dim=0).tolist() 49 | message += "The titles of each paper:\n" 50 | for i in range(batch.x.shape[0]): 51 | message += "node " + str(i) + " is '" + titles[i][0] + "'\n" 52 | message += "The sampled subgraph of the network is " + str(graph) + ' where the first number indicates source node and the second destination node.\n' 53 | message += "Hint: the direction of the edge can indicate some information of temporality.\n" 54 | message += "\nAccording to principles of citation network construction and the given subgraph structure, answer the following questions:\n" 55 | 56 | # In batch.edge_label_index.T.tolist(), index 0 and 1 are positive samples, 57 | # index 2 and 3 are negative samples. 58 | message += "Question 1: predict whether there tends to form an edge "+str(batch.edge_label_index.T.tolist()[1])+".\n" 59 | message += "Question 2: predict whether there tends to form an edge "+str(batch.edge_label_index.T.tolist()[3])+".\n" 60 | message += "Question 3: predict whether there tends to form an edge "+str(batch.edge_label_index.T.tolist()[2])+".\n" 61 | message += "Question 4: predict whether there tends to form an edge "+str(batch.edge_label_index.T.tolist()[0])+".\n" 62 | if reason: 63 | message += "Answer yes or no and show reasoning process.\n\n" 64 | else: 65 | message += "Answer yes or no and don't show any reasoning process.\n\n" 66 | 67 | return message 68 | -------------------------------------------------------------------------------- /examples/igbh/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-devel 2 | 3 | WORKDIR /workspace/repository 4 | 5 | RUN pip install torch==1.13.0+cu117 torchvision==0.14.0+cu117 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117 6 | RUN pip install scikit-learn==0.24.2 7 | RUN pip install torch_geometric==2.4.0 8 | RUN pip install --no-index torch_scatter==2.1.1 torch_sparse==0.6.17 -f https://data.pyg.org/whl/torch-1.13.0+cu117.html 9 | RUN pip install graphlearn-torch==0.2.2 10 | 11 | RUN apt update 12 | RUN apt install -y git 13 | RUN pip install git+https://github.com/mlcommons/logging.git 14 | 15 | # TF32 instead of FP32 for faster compute 16 | ENV NVIDIA_TF32_OVERRIDE=1 17 | 18 | RUN git clone https://github.com/alibaba/graphlearn-for-pytorch.git 19 | WORKDIR /workspace/repository/graphlearn-for-pytorch/examples/igbh 20 | -------------------------------------------------------------------------------- /examples/igbh/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alibaba/graphlearn-for-pytorch/26fe3d4e050b081bc51a79dc9547f244f5d314da/examples/igbh/__init__.py -------------------------------------------------------------------------------- /examples/igbh/build_partition_feature.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import argparse 17 | import os.path as osp 18 | 19 | import graphlearn_torch as glt 20 | import torch 21 | 22 | from dataset import IGBHeteroDataset 23 | 24 | 25 | def partition_feature(src_path: str, 26 | dst_path: str, 27 | partition_idx: int, 28 | chunk_size: int, 29 | dataset_size: str='tiny', 30 | in_memory: bool=True, 31 | use_fp16: bool=False): 32 | print(f'-- Loading igbh_{dataset_size} ...') 33 | data = IGBHeteroDataset(src_path, dataset_size, in_memory, with_edges=False, use_fp16=use_fp16) 34 | 35 | print(f'-- Build feature for partition {partition_idx} ...') 36 | dst_path = osp.join(dst_path, f'{dataset_size}-partitions') 37 | node_feat_dtype = torch.float16 if use_fp16 else torch.float32 38 | glt.partition.base.build_partition_feature(root_dir = dst_path, 39 | partition_idx = partition_idx, 40 | chunk_size = chunk_size, 41 | node_feat = data.feat_dict, 42 | node_feat_dtype = node_feat_dtype) 43 | 44 | 45 | if __name__ == '__main__': 46 | root = osp.join(osp.dirname(osp.dirname(osp.dirname(osp.realpath(__file__)))), 'data', 'igbh') 47 | glt.utils.ensure_dir(root) 48 | parser = argparse.ArgumentParser(description="Arguments for partitioning ogbn datasets.") 49 | parser.add_argument('--src_path', type=str, default=root, 50 | help='path containing the datasets') 51 | parser.add_argument('--dst_path', type=str, default=root, 52 | help='path containing the partitioned datasets') 53 | parser.add_argument('--dataset_size', type=str, default='full', 54 | choices=['tiny', 'small', 'medium', 'large', 'full'], 55 | help='size of the datasets') 56 | parser.add_argument('--in_memory', type=int, default=0, 57 | choices=[0, 1], help='0:read only mmap_mode=r, 1:load into memory') 58 | parser.add_argument("--partition_idx", type=int, default=0, 59 | help="Index of a partition") 60 | parser.add_argument("--chunk_size", type=int, default=10000, 61 | help="Chunk size for feature partitioning.") 62 | parser.add_argument("--use_fp16", action="store_true", 63 | help="save node/edge feature using fp16 format") 64 | 65 | 66 | args = parser.parse_args() 67 | 68 | partition_feature( 69 | args.src_path, 70 | args.dst_path, 71 | partition_idx=args.partition_idx, 72 | chunk_size=args.chunk_size, 73 | dataset_size=args.dataset_size, 74 | in_memory=args.in_memory==1, 75 | use_fp16=args.use_fp16 76 | ) 77 | -------------------------------------------------------------------------------- /examples/igbh/download.py: -------------------------------------------------------------------------------- 1 | import tarfile, hashlib, os 2 | import os.path as osp 3 | from tqdm import tqdm 4 | import urllib.request as ur 5 | 6 | # https://github.com/IllinoisGraphBenchmark/IGB-Datasets/blob/main/igb/download.py 7 | 8 | GBFACTOR = float(1 << 30) 9 | 10 | def decide_download(url): 11 | d = ur.urlopen(url) 12 | size = int(d.info()["Content-Length"])/GBFACTOR 13 | ### confirm if larger than 1GB 14 | if size > 1: 15 | return input("This will download %.2fGB. Will you proceed? (y/N) " % (size)).lower() == "y" 16 | else: 17 | return True 18 | 19 | 20 | dataset_urls = { 21 | 'homogeneous' : { 22 | 'tiny' : 'https://igb-public.s3.us-east-2.amazonaws.com/igb-homogeneous/igb_homogeneous_tiny.tar.gz', 23 | 'small' : 'https://igb-public.s3.us-east-2.amazonaws.com/igb-homogeneous/igb_homogeneous_small.tar.gz', 24 | 'medium' : 'https://igb-public.s3.us-east-2.amazonaws.com/igb-homogeneous/igb_homogeneous_medium.tar.gz' 25 | }, 26 | 'heterogeneous' : { 27 | 'tiny' : 'https://igb-public.s3.us-east-2.amazonaws.com/igb-heterogeneous/igb_heterogeneous_tiny.tar.gz', 28 | 'small' : 'https://igb-public.s3.us-east-2.amazonaws.com/igb-heterogeneous/igb_heterogeneous_small.tar.gz', 29 | 'medium' : 'https://igb-public.s3.us-east-2.amazonaws.com/igb-heterogeneous/igb_heterogeneous_medium.tar.gz' 30 | } 31 | } 32 | 33 | 34 | md5checksums = { 35 | 'homogeneous' : { 36 | 'tiny' : '34856534da55419b316d620e2d5b21be', 37 | 'small' : '6781c699723529902ace0a95cafe6fe4', 38 | 'medium' : '4640df4ceee46851fd18c0a44ddcc622' 39 | }, 40 | 'heterogeneous' : { 41 | 'tiny' : '83fbc1091497ff92cf20afe82fae0ade', 42 | 'small' : '2f42077be60a074aec24f7c60089e1bd', 43 | 'medium' : '7f0df4296eca36553ff3a6a63abbd347' 44 | } 45 | } 46 | 47 | 48 | def check_md5sum(dataset_type, dataset_size, filename): 49 | original_md5 = md5checksums[dataset_type][dataset_size] 50 | 51 | with open(filename, 'rb') as file_to_check: 52 | data = file_to_check.read() 53 | md5_returned = hashlib.md5(data).hexdigest() 54 | 55 | if original_md5 == md5_returned: 56 | print(" md5sum verified.") 57 | return 58 | else: 59 | os.remove(filename) 60 | raise Exception(" md5sum verification failed!.") 61 | 62 | 63 | def download_dataset(path, dataset_type, dataset_size): 64 | output_directory = path 65 | url = dataset_urls[dataset_type][dataset_size] 66 | if decide_download(url): 67 | data = ur.urlopen(url) 68 | size = int(data.info()["Content-Length"]) 69 | chunk_size = 1024*1024 70 | num_iter = int(size/chunk_size) + 2 71 | downloaded_size = 0 72 | filename = path + "/igb_" + dataset_type + "_" + dataset_size + ".tar.gz" 73 | with open(filename, 'wb') as f: 74 | pbar = tqdm(range(num_iter)) 75 | for _ in pbar: 76 | chunk = data.read(chunk_size) 77 | downloaded_size += len(chunk) 78 | pbar.set_description("Downloaded {:.2f} GB".format(float(downloaded_size)/GBFACTOR)) 79 | f.write(chunk) 80 | print("Downloaded" + " igb_" + dataset_type + "_" + dataset_size, end=" ->") 81 | check_md5sum(dataset_type, dataset_size, filename) 82 | file = tarfile.open(filename) 83 | file.extractall(output_directory) 84 | file.close() 85 | size = 0 86 | for path, _, files in os.walk(output_directory+"/"+dataset_size): 87 | for f in files: 88 | fp = osp.join(path, f) 89 | size += osp.getsize(fp) 90 | print("Final dataset size {:.2f} GB.".format(size/GBFACTOR)) 91 | os.remove(filename) 92 | -------------------------------------------------------------------------------- /examples/igbh/download_igbh_full.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | #https://github.com/IllinoisGraphBenchmark/IGB-Datasets/blob/main/igb/download_igbh600m.sh 4 | echo "IGBH600M download starting" 5 | cd ../../data/ 6 | mkdir -p igbh/full/processed 7 | cd igbh/full/processed 8 | 9 | # paper 10 | mkdir paper 11 | cd paper 12 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper/node_feat.npy 13 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper/node_label_19.npy 14 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper/node_label_2K.npy 15 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper/paper_id_index_mapping.npy 16 | cd .. 17 | 18 | # paper__cites__paper 19 | mkdir paper__cites__paper 20 | cd paper__cites__paper 21 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper__cites__paper/edge_index.npy 22 | cd .. 23 | 24 | # author 25 | mkdir author 26 | cd author 27 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/author/author_id_index_mapping.npy 28 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/author/node_feat.npy 29 | cd .. 30 | 31 | # conference 32 | mkdir conference 33 | cd conference 34 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/conference/conference_id_index_mapping.npy 35 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/conference/node_feat.npy 36 | cd .. 37 | 38 | # institute 39 | mkdir institute 40 | cd institute 41 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/institute/institute_id_index_mapping.npy 42 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/institute/node_feat.npy 43 | cd .. 44 | 45 | # journal 46 | mkdir journal 47 | cd journal 48 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/journal/journal_id_index_mapping.npy 49 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/journal/node_feat.npy 50 | cd .. 51 | 52 | # fos 53 | mkdir fos 54 | cd fos 55 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/fos/fos_id_index_mapping.npy 56 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/fos/node_feat.npy 57 | cd .. 58 | 59 | # author__affiliated_to__institute 60 | mkdir author__affiliated_to__institute 61 | cd author__affiliated_to__institute 62 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/author__affiliated_to__institute/edge_index.npy 63 | cd .. 64 | 65 | # paper__published__journal 66 | mkdir paper__published__journal 67 | cd paper__published__journal 68 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper__published__journal/edge_index.npy 69 | cd .. 70 | 71 | # paper__topic__fos 72 | mkdir paper__topic__fos 73 | cd paper__topic__fos 74 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper__topic__fos/edge_index.npy 75 | cd .. 76 | 77 | # paper__venue__conference 78 | mkdir paper__venue__conference 79 | cd paper__venue__conference 80 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper__venue__conference/edge_index.npy 81 | cd .. 82 | 83 | # paper__written_by__author 84 | mkdir paper__written_by__author 85 | cd paper__written_by__author 86 | wget -c https://igb-public.s3.us-east-2.amazonaws.com/IGBH/processed/paper__written_by__author/edge_index.npy 87 | cd .. 88 | 89 | echo "IGBH-IGBH download complete" 90 | -------------------------------------------------------------------------------- /examples/igbh/mlperf_logging_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | from mlperf_logging import mllog 3 | from mlperf_logging.mllog import constants 4 | from mlperf_logging.mllog.mllog import MLLogger 5 | 6 | def get_mlperf_logger(path, filename='mlperf_gnn.log'): 7 | mllog.config(filename=os.path.join(path, filename)) 8 | mllogger = mllog.get_mllogger() 9 | mllogger.logger.propagate = False 10 | return mllogger 11 | 12 | def submission_info(mllogger: MLLogger, benchmark_name: str, submitter_name: str): 13 | """Logs required for a valid MLPerf submission.""" 14 | mllogger.event( 15 | key=constants.SUBMISSION_BENCHMARK, 16 | value=benchmark_name, 17 | ) 18 | mllogger.event( 19 | key=constants.SUBMISSION_ORG, 20 | value=submitter_name, 21 | ) 22 | mllogger.event( 23 | key=constants.SUBMISSION_DIVISION, 24 | value=constants.CLOSED, 25 | ) 26 | mllogger.event( 27 | key=constants.SUBMISSION_STATUS, 28 | value=constants.ONPREM, 29 | ) 30 | mllogger.event( 31 | key=constants.SUBMISSION_PLATFORM, 32 | value=submitter_name, 33 | ) 34 | -------------------------------------------------------------------------------- /examples/igbh/rgnn.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import torch 17 | import torch.nn.functional as F 18 | 19 | from torch_geometric.nn import HeteroConv, GATConv, GCNConv, SAGEConv 20 | from torch_geometric.utils import trim_to_layer 21 | 22 | class RGNN(torch.nn.Module): 23 | r""" [Relational GNN model](https://arxiv.org/abs/1703.06103). 24 | 25 | Args: 26 | etypes: edge types. 27 | in_dim: input size. 28 | h_dim: Dimension of hidden layer. 29 | out_dim: Output dimension. 30 | num_layers: Number of conv layers. 31 | dropout: Dropout probability for hidden layers. 32 | model: "rsage" or "rgat". 33 | heads: Number of multi-head-attentions for GAT. 34 | node_type: The predict node type for node classification. 35 | 36 | """ 37 | def __init__(self, etypes, in_dim, h_dim, out_dim, num_layers=2, 38 | dropout=0.2, model='rgat', heads=4, node_type=None, with_trim=False): 39 | super().__init__() 40 | self.node_type = node_type 41 | if node_type is not None: 42 | self.lin = torch.nn.Linear(h_dim, out_dim) 43 | 44 | self.convs = torch.nn.ModuleList() 45 | for i in range(num_layers): 46 | in_dim = in_dim if i == 0 else h_dim 47 | h_dim = out_dim if (i == (num_layers - 1) and node_type is None) else h_dim 48 | if model == 'rsage': 49 | self.convs.append(HeteroConv({ 50 | etype: SAGEConv(in_dim, h_dim, root_weight=False) 51 | for etype in etypes})) 52 | elif model == 'rgat': 53 | self.convs.append(HeteroConv({ 54 | etype: GATConv(in_dim, h_dim // heads, heads=heads, add_self_loops=False) 55 | for etype in etypes})) 56 | self.dropout = torch.nn.Dropout(dropout) 57 | self.with_trim = with_trim 58 | 59 | def forward(self, x_dict, edge_index_dict, num_sampled_edges_dict=None, 60 | num_sampled_nodes_dict=None): 61 | for i, conv in enumerate(self.convs): 62 | if self.with_trim: 63 | x_dict, edge_index_dict, _ = trim_to_layer( 64 | layer=i, 65 | num_sampled_nodes_per_hop=num_sampled_nodes_dict, 66 | num_sampled_edges_per_hop=num_sampled_edges_dict, 67 | x=x_dict, 68 | edge_index=edge_index_dict 69 | ) 70 | for key in list(edge_index_dict.keys()): 71 | if key[0] not in x_dict or key[-1] not in x_dict: 72 | del edge_index_dict[key] 73 | 74 | x_dict = conv(x_dict, edge_index_dict) 75 | if i != len(self.convs) - 1: 76 | x_dict = {key: F.leaky_relu(x) for key, x in x_dict.items()} 77 | x_dict = {key: self.dropout(x) for key, x in x_dict.items()} 78 | if hasattr(self, 'lin'): # for node classification 79 | return self.lin(x_dict[self.node_type]) 80 | else: 81 | return x_dict 82 | -------------------------------------------------------------------------------- /examples/igbh/split_seeds.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os.path as osp 3 | 4 | import torch 5 | 6 | class SeedSplitter(object): 7 | def __init__(self, 8 | path, 9 | dataset_size='tiny', 10 | use_label_2K=True, 11 | random_seed=42, 12 | validation_frac=0.01): 13 | self.path = path 14 | self.dataset_size = dataset_size 15 | self.use_label_2K = use_label_2K 16 | self.random_seed = random_seed 17 | self.validation_frac = validation_frac 18 | self.paper_nodes_num = {'tiny':100000, 'small':1000000, 'medium':10000000, 'large':100000000, 'full':269346174} 19 | self.process() 20 | 21 | def process(self): 22 | torch.manual_seed(self.random_seed) 23 | n_labeled_idx = self.paper_nodes_num[self.dataset_size] 24 | if self.dataset_size == 'full': 25 | if self.use_label_2K: 26 | n_labeled_idx = 157675969 27 | else: 28 | n_labeled_idx = 227130858 29 | 30 | shuffled_index = torch.randperm(n_labeled_idx) 31 | n_train = int(n_labeled_idx * 0.6) 32 | n_val = int(n_labeled_idx * self.validation_frac) 33 | 34 | train_idx = shuffled_index[:n_train] 35 | val_idx = shuffled_index[n_train : n_train + n_val] 36 | 37 | path = osp.join(self.path, self.dataset_size, 'processed') 38 | torch.save(train_idx, osp.join(path, 'train_idx.pt')) 39 | torch.save(val_idx, osp.join(path, 'val_idx.pt')) 40 | 41 | if __name__ == '__main__': 42 | parser = argparse.ArgumentParser() 43 | root = osp.join(osp.dirname(osp.dirname(osp.dirname(osp.realpath(__file__)))), 'data', 'igbh') 44 | parser.add_argument('--path', type=str, default=root, 45 | help='path containing the datasets') 46 | parser.add_argument('--dataset_size', type=str, default='full', 47 | choices=['tiny', 'small', 'medium', 'large', 'full'], 48 | help='size of the datasets') 49 | parser.add_argument("--random_seed", type=int, default='42') 50 | parser.add_argument('--num_classes', type=int, default=2983, 51 | choices=[19, 2983], help='number of classes') 52 | parser.add_argument("--validation_frac", type=float, default=0.005, 53 | help="Fraction of labeled vertices to be used for validation.") 54 | 55 | args = parser.parse_args() 56 | splitter = SeedSplitter(path=args.path, 57 | dataset_size=args.dataset_size, 58 | use_label_2K=(args.num_classes==2983), 59 | random_seed=args.random_seed, 60 | validation_frac=args.validation_frac) -------------------------------------------------------------------------------- /examples/igbh/utilities.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import torch 4 | 5 | def create_ckpt_folder(base_dir, prefix="ckpt"): 6 | timestamp = time.strftime("%Y%m%d-%H%M%S") 7 | folder_name = f"{prefix}_{timestamp}" if prefix else timestamp 8 | full_path = os.path.join(base_dir, folder_name) 9 | if not os.path.exists(full_path): 10 | os.makedirs(full_path) 11 | return full_path 12 | 13 | -------------------------------------------------------------------------------- /examples/pai/README.md: -------------------------------------------------------------------------------- 1 | ## Run on PAI with MaxCompute(ODPS) table as input. 2 | 3 | - ogbn_products: Node classification examples with ogbn_products. -------------------------------------------------------------------------------- /examples/pai/ogbn_products/README.md: -------------------------------------------------------------------------------- 1 | ## Run on PAI with MaxCompute(ODPS) table as input. 2 | 3 | ### 1. Prepare files 4 | ``` 5 | cd .. 6 | tar -zcvf train_products_sage.tar.gz ogbn_products/*.py requirements.txt 7 | ``` 8 | 9 | 10 | ### 2. Input 11 | ``` 12 | ogbn_products_node(id bigint, feature string) 13 | ogbn_products_edge(src_id bigint, dst_id bigint, e_id bigint) 14 | ogbn_products_train(id bigint, label bigint) 15 | ``` 16 | You can first generate text files use 17 | ``` 18 | python data_preprocess.py 19 | ``` 20 | 21 | and then use tunnel to upload text files to ODPS Tables 22 | ``` 23 | create table if not exists ogbn_products_node(id bigint, feature string); 24 | tunnel u /{your_path}/graphlearn-for-pytorch/data/products/ogbn_products_node ogbn_products_node -fd "\t"; 25 | create table if not exists ogbn_products_edge(src_id bigint, dst_id bigint, e_id bigint); 26 | tunnel u /{your_path}/graphlearn-for-pytorch/data/products/ogbn_products_edge ogbn_products_edge -fd "\t"; 27 | 28 | create table if not exists ogbn_products_train(id bigint, label bigint); 29 | tunnel u /{your_path}/graphlearn-for-pytorch/data/products/ogbn_products_train ogbn_products_train -fd "\t"; 30 | ``` 31 | 32 | Note: **The src_id and dst_id of the edge and the id of the node must correspond, 33 | and the ids are required to be encoded consecutively starting from 0.** 34 | 35 | 36 | ### 3. PAI command 37 | #### 3.1 Single node single GPU 38 | 39 | ``` 40 | pai -name pytorch112z -Dscript='file://{your_path}/graphlearn-for-pytorch/examples/pai/train_products_sage.tar.gz' -DentryFile='ogbn_products/train_products_sage.py' -Dtables="odps://{your_project}/tables/ogbn_products_node,odps://{your_project}/tables/ogbn_products_edge,odps://{your_project}/tables/ogbn_products_train" -Dcluster="{\"worker\":{\"gpu\":100}}" -DworkerCount=1 -DuserDefinedParameters='--split_ratio=0.2'; 41 | ``` 42 | 43 | #### 3.2 Single node(machine) multi GPUs 44 | 1 node with 4 GPUs 45 | 46 | ``` 47 | pai -name pytorch112z -Dscript='file://{your_path}/graphlearn-for-pytorch/examples/pai/train_products_sage.tar.gz' -DentryFile='ogbn_products/train_products_sage.py' -Dtables="odps://{your_project}/tables/ogbn_products_node,odps://{your_project}/tables/ogbn_products_edge,odps://{your_project}/tables/ogbn_products_train" -Dcluster="{\"worker\":{\"gpu\":400}}" -DworkerCount=1 -DuserDefinedParameters='--split_ratio=0.2'; 48 | ``` 49 | 50 | #### 3.3 Mulit-nodes mulit GPUs 51 | 2 nodes each with 2 GPUs 52 | 53 | ``` 54 | pai -name pytorch112z -Dscript='file:///{your_path}/graphlearn-for-pytorch/examples/pai/train_products_sage.tar.gz' -DentryFile='ogbn_products/dist_train_products_sage.py' -Dtables="odps://{your_project}/tables/ogbn_products_node,odps://{your_project}/tables/ogbn_products_edge,odps://{your_project}/tables/ogbn_products_train" -Dcluster="{\"worker\":{\"gpu\":200}}" -DworkerCount=2 -DuserDefinedParameters='--num_training_procs=2'; 55 | ``` 56 | -------------------------------------------------------------------------------- /examples/pai/ogbn_products/data_preprocess.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import os.path as osp 17 | from ogb.nodeproppred import PygNodePropPredDataset 18 | 19 | 20 | import numpy as np 21 | from ogb.nodeproppred import NodePropPredDataset 22 | 23 | # load data 24 | root = osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), 25 | '../..', 'data', 'products') 26 | dataset = NodePropPredDataset('ogbn-products', root) 27 | split_idx = dataset.get_idx_split() 28 | train_idx = split_idx['train'] 29 | graph, label = dataset[0] # label with shape(2449029, 1) 30 | num_nodes = graph['num_nodes'] # 2449029 31 | node_feat = graph['node_feat'] # shape(2449029, 100) 32 | edge_index = graph['edge_index'] # shape(2, 123718280) 33 | 34 | # dump to disk 35 | train_table = osp.join(root, 'ogbn_products_train') 36 | node_table = osp.join(root, 'ogbn_products_node') 37 | edge_table = osp.join(root, 'ogbn_products_edge') 38 | 39 | with open(train_table, 'w') as f: 40 | for i in train_idx: 41 | f.write(str(i) + '\t' + str(label[i][0]) + '\n') 42 | 43 | with open(node_table, 'w') as f: 44 | for i in range(num_nodes): 45 | f.write(str(i) + '\t' + str(':'.join(map(str, node_feat[i]))) + '\n') 46 | 47 | with open(edge_table, 'w') as f: 48 | for i in range(edge_index.shape[1]): 49 | f.write(str(edge_index[0][i]) + '\t' + str(edge_index[1][i]) + '\t' + str(i) + '\n') -------------------------------------------------------------------------------- /examples/pai/requirements.txt: -------------------------------------------------------------------------------- 1 | https://graphlearn.oss-cn-hangzhou.aliyuncs.com/graphlearn_torch-0.2.0-cp38-cp38-linux_x86_64.whl -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/graph.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include "graphlearn_torch/include/graph.h" 17 | #include "graphlearn_torch/include/common.h" 18 | 19 | namespace graphlearn_torch { 20 | 21 | void Graph::InitCPUGraphFromCSR( 22 | const torch::Tensor& indptr, 23 | const torch::Tensor& indices, 24 | const torch::Tensor& edge_ids, 25 | const torch::Tensor& edge_weights) { 26 | CheckEq(indptr.dim(), 1); 27 | CheckEq(indices.dim(), 1); 28 | 29 | row_ptr_ = indptr.data_ptr(); 30 | col_idx_ = indices.data_ptr(); 31 | row_count_ = indptr.size(0) - 1; 32 | edge_count_ = indices.size(0); 33 | col_count_ = std::get<0>(at::_unique(indices)).size(0); 34 | 35 | if (edge_ids.numel()) { 36 | CheckEq(edge_ids.dim(), 1); 37 | CheckEq(edge_ids.numel(), indices.numel()); 38 | edge_id_ = edge_ids.data_ptr(); 39 | } 40 | 41 | if (edge_weights.numel()) { 42 | CheckEq(edge_weights.dim(), 1); 43 | CheckEq(edge_weights.numel(), indices.numel()); 44 | edge_weight_ = edge_weights.data_ptr(); 45 | } 46 | } 47 | 48 | } // namespace graphlearn_torch 49 | -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/inducer.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CPU_INDUCER_H_ 17 | #define GRAPHLEARN_TORCH_CPU_INDUCER_H_ 18 | 19 | #include 20 | #include 21 | #include 22 | 23 | #include "graphlearn_torch/include/inducer_base.h" 24 | #include "graphlearn_torch/include/types.h" 25 | 26 | namespace graphlearn_torch { 27 | 28 | class CPUInducer : public Inducer { 29 | public: 30 | explicit CPUInducer(int32_t num_nodes); 31 | virtual ~CPUInducer() {} 32 | CPUInducer(const CPUInducer& other) = delete; 33 | CPUInducer& operator=(const CPUInducer& other) = delete; 34 | CPUInducer(CPUInducer&&) = default; 35 | CPUInducer& operator=(CPUInducer&&) = default; 36 | 37 | virtual torch::Tensor InitNode(const torch::Tensor& seed); 38 | virtual std::tuple 39 | InduceNext(const torch::Tensor& srcs, 40 | const torch::Tensor& nbrs, 41 | const torch::Tensor& nbrs_num); 42 | virtual void Reset(); 43 | 44 | private: 45 | int32_t nodes_size_; 46 | IntHashMap glob2local_; 47 | }; 48 | 49 | 50 | class CPUHeteroInducer : public HeteroInducer{ 51 | public: 52 | explicit CPUHeteroInducer(std::unordered_map num_nodes); 53 | virtual ~CPUHeteroInducer() {} 54 | CPUHeteroInducer(const CPUHeteroInducer& other) = delete; 55 | CPUHeteroInducer& operator=(const CPUHeteroInducer& other) = delete; 56 | CPUHeteroInducer(CPUHeteroInducer&&) = default; 57 | CPUHeteroInducer& operator=(CPUHeteroInducer&&) = default; 58 | 59 | virtual TensorMap InitNode(const TensorMap& seed); 60 | virtual HeteroCOO InduceNext(const HeteroNbr& nbrs); 61 | virtual void Reset(); 62 | 63 | private: 64 | void InsertGlob2Local(const HeteroNbr& nbrs, 65 | std::unordered_map>& out_nodes); 66 | 67 | private: 68 | std::unordered_map nodes_size_; 69 | std::unordered_map glob2local_; 70 | }; 71 | 72 | } // namespace graphlearn_torch 73 | 74 | #endif // GRAPHLEARN_TORCH_CPU_INDUCER_H_ 75 | -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/random_negative_sampler.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include "graphlearn_torch/include/common.h" 17 | #include "graphlearn_torch/csrc/cpu/random_negative_sampler.h" 18 | 19 | #include 20 | 21 | 22 | namespace graphlearn_torch { 23 | 24 | std::tuple CPURandomNegativeSampler::Sample( 25 | int32_t req_num, int32_t trials_num, bool padding) { 26 | const int64_t* row_ptr = graph_->GetRowPtr(); 27 | const int64_t* col_idx = graph_->GetColIdx(); 28 | int64_t row_num = graph_->GetRowCount(); 29 | int64_t col_num = graph_->GetColCount(); 30 | uint32_t seed = RandomSeedManager::getInstance().getSeed(); 31 | thread_local static std::mt19937 engine(seed); 32 | std::uniform_int_distribution row_dist(0, row_num - 1); 33 | std::uniform_int_distribution col_dist(0, col_num - 1); 34 | int64_t row_data[req_num]; 35 | int64_t col_data[req_num]; 36 | int32_t out_prefix[req_num]; 37 | std::fill(out_prefix, out_prefix + req_num, 0); 38 | 39 | at::parallel_for(0, req_num, 1, [&](int32_t start, int32_t end) { 40 | for(int32_t i = start; i < end; ++i) { 41 | for(int32_t j = 0; j < trials_num; ++j) { 42 | int64_t r = row_dist(engine); 43 | int64_t c = col_dist(engine); 44 | if (!EdgeInCSR(row_ptr, col_idx, r, c)) { 45 | row_data[i] = r; 46 | col_data[i] = c; 47 | out_prefix[i] = 1; 48 | break; 49 | } 50 | } 51 | } 52 | }); 53 | // sort sampled results. 54 | int32_t cursor = 0; 55 | for (int32_t i = 0; i < req_num; ++i) { 56 | if (out_prefix[i] == 1) { 57 | row_data[cursor] = row_data[i]; 58 | col_data[cursor] = col_data[i]; 59 | ++cursor; 60 | } 61 | } 62 | int32_t sampled_num = std::accumulate(out_prefix, out_prefix + req_num, 0); 63 | while ((sampled_num < req_num) && padding) { // non-strict negative sampling. 64 | row_data[sampled_num] = row_dist(engine); 65 | col_data[sampled_num] = col_dist(engine); 66 | ++sampled_num; 67 | } 68 | 69 | torch::Tensor rows = torch::empty(sampled_num, torch::kInt64); 70 | torch::Tensor cols = torch::empty(sampled_num, torch::kInt64); 71 | std::copy(row_data, row_data + sampled_num, rows.data_ptr()); 72 | std::copy(col_data, col_data + sampled_num, cols.data_ptr()); 73 | return std::make_tuple(rows, cols); 74 | } 75 | 76 | bool CPURandomNegativeSampler::EdgeInCSR(const int64_t* row_ptr, 77 | const int64_t* col_idx, 78 | int64_t r, 79 | int64_t c) { 80 | const int64_t* start = col_idx + row_ptr[r]; 81 | const int64_t* end = col_idx + row_ptr[r + 1]; 82 | return std::binary_search(start, end, c); 83 | } 84 | 85 | } // namespace graphlearn_torch 86 | -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/random_negative_sampler.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CPU_RANDOM_NEGATIVE_SAMPLER_H_ 17 | #define GRAPHLEARN_TORCH_CPU_RANDOM_NEGATIVE_SAMPLER_H_ 18 | 19 | #include "graphlearn_torch/include/negative_sampler.h" 20 | 21 | namespace graphlearn_torch { 22 | 23 | class CPURandomNegativeSampler : public NegativeSampler { 24 | public: 25 | CPURandomNegativeSampler(const Graph* graph) : NegativeSampler(graph) {} 26 | ~CPURandomNegativeSampler() {} 27 | virtual std::tuple Sample( 28 | int32_t req_num, int32_t trials_num, bool padding=false); 29 | 30 | private: 31 | bool EdgeInCSR(const int64_t* row_ptr, 32 | const int64_t* col_idx, 33 | int64_t r, 34 | int64_t c); 35 | }; 36 | 37 | } // namespace graphlearn_torch 38 | 39 | #endif // GRAPHLEARN_TORCH_CPU_RANDOM_NEGATIVE_SAMPLER_H_ -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/random_sampler.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CPU_RANDOM_SAMPLER_H_ 17 | #define GRAPHLEARN_TORCH_CPU_RANDOM_SAMPLER_H_ 18 | 19 | #include "graphlearn_torch/include/sampler.h" 20 | 21 | 22 | namespace graphlearn_torch { 23 | 24 | class CPURandomSampler : public Sampler { 25 | public: 26 | CPURandomSampler(const Graph* graph) : Sampler(graph) {} 27 | ~CPURandomSampler() {} 28 | 29 | std::tuple Sample( 30 | const torch::Tensor& nodes, int32_t req_num) override; 31 | 32 | std::tuple 33 | SampleWithEdge(const torch::Tensor& nodes, int32_t req_num) override; 34 | 35 | private: 36 | void FillNbrsNum(const int64_t* nodes, const int32_t bs, 37 | const int32_t req_num, const int64_t row_count, 38 | const int64_t* row_ptr, int64_t* out_nbr_num); 39 | 40 | void CSRRowWiseSample(const int64_t* nodes, const int64_t* nbrs_offset, 41 | const int32_t bs, const int32_t req_num, const int64_t row_count, 42 | const int64_t* row_ptr, const int64_t* col_idx, int64_t* out_nbrs); 43 | 44 | void CSRRowWiseSample(const int64_t* nodes, const int64_t* nbrs_offset, 45 | const int32_t bs, const int32_t req_num, const int64_t row_count, 46 | const int64_t* row_ptr, const int64_t* col_idx, const int64_t* edge_id, 47 | int64_t* out_nbrs, int64_t* out_eid); 48 | 49 | void UniformSample(const int64_t* col_begin, const int64_t* col_end, 50 | const int32_t req_num, int64_t* out_nbrs); 51 | 52 | void UniformSample(const int64_t* col_begin, const int64_t* col_end, 53 | const int64_t* eid_begin, const int64_t* eid_end, 54 | const int32_t req_num, int64_t* out_nbrs, int64_t* out_eid); 55 | 56 | }; 57 | 58 | } // namespace graphlearn_torch 59 | 60 | #endif // GRAPHLEARN_TORCH_CPU_RANDOM_SAMPLER_H_ -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/stitch_sample_results.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include "graphlearn_torch/include/stitch_sample_results.h" 17 | 18 | namespace graphlearn_torch { 19 | 20 | std::tuple> 21 | CPUStitchSampleResults(const torch::Tensor& ids, 22 | const std::vector& idx_list, 23 | const std::vector& nbrs_list, 24 | const std::vector& nbrs_num_list, 25 | const std::vector& eids_list) { 26 | int64_t ids_count = ids.size(0); 27 | int64_t partitions = idx_list.size(); 28 | bool with_edge = !eids_list.empty(); 29 | auto options = torch::TensorOptions() 30 | .dtype(torch::kInt64) 31 | .device(ids.device()); 32 | 33 | auto nbrs_num = torch::zeros(ids_count, options); 34 | for (int64_t i = 0; i < partitions; ++i) { 35 | nbrs_num = nbrs_num.index_copy(0, idx_list[i], nbrs_num_list[i]); 36 | } 37 | auto nbrs_offsets = torch::cumsum(nbrs_num, 0, torch::kInt64); 38 | const int64_t* nbrs_offset_data = nbrs_offsets.data_ptr(); 39 | auto nbrs_count = nbrs_offset_data[ids_count - 1]; 40 | auto nbrs = torch::zeros(nbrs_count, options); 41 | int64_t* nbrs_data = nbrs.data_ptr(); 42 | auto eids = with_edge ? 43 | torch::zeros(nbrs_count, options) : torch::empty(0, options); 44 | int64_t* eids_data = with_edge ? eids.data_ptr() : nullptr; 45 | 46 | for (int64_t i = 0; i < partitions; ++i) { 47 | int64_t partital_count = idx_list[i].size(0); 48 | auto p_nbrs_offsets = torch::cumsum(nbrs_num_list[i], 0); 49 | const int64_t* p_nbrs_offsets_data = p_nbrs_offsets.data_ptr(); 50 | const int64_t* p_idx_data = idx_list[i].data_ptr(); 51 | const int64_t* p_nbrs_data = nbrs_list[i].data_ptr(); 52 | const int64_t* p_nbrs_num_data = nbrs_num_list[i].data_ptr(); 53 | const int64_t* p_eids_data = with_edge ? 54 | eids_list[i].data_ptr() : nullptr; 55 | at::parallel_for(0, partital_count, 1, 56 | [p_idx_data, p_nbrs_data, p_nbrs_num_data, p_eids_data, 57 | p_nbrs_offsets_data, nbrs_offset_data, with_edge, nbrs_data, eids_data] 58 | (int32_t start, int32_t end) { 59 | for (int32_t i = start; i < end; i++) { 60 | auto cur_id_idx = p_idx_data[i]; 61 | auto cur_nbr_num = p_nbrs_num_data[i]; 62 | if (cur_nbr_num > 0) { 63 | auto local_offset = p_nbrs_offsets_data[i] - cur_nbr_num; 64 | auto global_offset = nbrs_offset_data[cur_id_idx] - cur_nbr_num; 65 | memcpy(nbrs_data + global_offset, 66 | p_nbrs_data + local_offset, 67 | sizeof(int64_t) * cur_nbr_num); 68 | if (with_edge) { 69 | memcpy(eids_data + global_offset, 70 | p_eids_data + local_offset, 71 | sizeof(int64_t) * cur_nbr_num); 72 | } 73 | } 74 | } 75 | }); 76 | } 77 | 78 | return std::make_tuple( 79 | nbrs, 80 | nbrs_num, 81 | with_edge ? torch::optional{eids} : 82 | torch::optional{torch::nullopt}); 83 | } 84 | 85 | } // namespace graphlearn_torch 86 | -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/subgraph_op.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include "graphlearn_torch/csrc/cpu/subgraph_op.h" 17 | 18 | namespace graphlearn_torch { 19 | 20 | SubGraph CPUSubGraphOp::NodeSubGraph(const torch::Tensor& srcs, 21 | bool with_edge) { 22 | std::vector out_nodes; 23 | InitNode(srcs, out_nodes); 24 | torch::Tensor nodes = torch::empty(out_nodes.size(), srcs.options()); 25 | std::copy(out_nodes.begin(), out_nodes.end(), nodes.data_ptr()); 26 | std::vector out_rows; 27 | std::vector out_cols; 28 | std::vector out_eids; 29 | Induce(out_nodes, with_edge, out_rows, out_cols, out_eids); 30 | 31 | auto edge_size = out_rows.size(); 32 | torch::Tensor rows = torch::empty(edge_size, srcs.options()); 33 | torch::Tensor cols = torch::empty(edge_size, srcs.options()); 34 | std::copy(out_rows.begin(), out_rows.end(), rows.data_ptr()); 35 | std::copy(out_cols.begin(), out_cols.end(), cols.data_ptr()); 36 | auto subgraph = SubGraph(nodes, rows, cols); 37 | if (with_edge) { 38 | torch::Tensor eids = torch::empty(edge_size, srcs.options()); 39 | std::copy(out_eids.begin(), out_eids.end(), eids.data_ptr()); 40 | subgraph.eids = eids; 41 | } 42 | 43 | return subgraph; 44 | } 45 | 46 | void CPUSubGraphOp::InitNode(const torch::Tensor& srcs, 47 | std::vector& out_nodes) { 48 | Reset(); 49 | const auto src_size = srcs.size(0); 50 | const auto src_ptr = srcs.data_ptr(); 51 | int32_t nodes_size = 0; 52 | out_nodes.reserve(src_size); 53 | for (int32_t i = 0; i < src_size; ++i) { 54 | if (glob2local_.insert(std::make_pair(src_ptr[i], nodes_size)).second) { 55 | out_nodes.push_back(src_ptr[i]); 56 | ++nodes_size; 57 | } 58 | } 59 | } 60 | 61 | void CPUSubGraphOp::Induce(const std::vector& nodes, 62 | bool with_edge, 63 | std::vector& out_rows, 64 | std::vector& out_cols, 65 | std::vector& out_eids) { 66 | const auto indptr = graph_->GetRowPtr(); 67 | const auto indices = graph_->GetColIdx(); 68 | const auto edge_ids = graph_->GetEdgeId(); 69 | const auto row_count = graph_->GetRowCount(); 70 | const auto node_size = nodes.size(); 71 | out_rows.reserve(node_size * node_size); 72 | out_cols.reserve(node_size * node_size); 73 | if (with_edge) out_eids.reserve(node_size * node_size); 74 | 75 | for (int32_t i = 0; i < node_size; ++i) { 76 | auto old_row = nodes[i]; 77 | if (old_row < row_count) { 78 | for (int32_t j = indptr[old_row]; j < indptr[old_row+1]; ++j) { 79 | auto old_col = indices[j]; 80 | auto col_iter = glob2local_.find(old_col); 81 | if (col_iter != glob2local_.end()) { 82 | out_rows.push_back(i); 83 | out_cols.push_back(col_iter->second); 84 | if (with_edge) out_eids.push_back(edge_ids[j]); 85 | } 86 | } 87 | } 88 | } 89 | } 90 | 91 | } // namespace graphlearn_torch -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/subgraph_op.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CSRC_CPU_SUBGRAPH_OP_H_ 17 | #define GRAPHLEARN_TORCH_CSRC_CPU_SUBGRAPH_OP_H_ 18 | 19 | #include "graphlearn_torch/include/subgraph_op_base.h" 20 | 21 | #include 22 | #include 23 | #include 24 | 25 | #include "graphlearn_torch/include/graph.h" 26 | #include "graphlearn_torch/include/types.h" 27 | 28 | namespace graphlearn_torch { 29 | 30 | class CPUSubGraphOp : public SubGraphOp { 31 | public: 32 | CPUSubGraphOp(const Graph* graph): SubGraphOp(graph) {} 33 | ~CPUSubGraphOp() {} 34 | SubGraph NodeSubGraph(const torch::Tensor& srcs, bool with_edge=false) override; 35 | 36 | private: 37 | void Reset() { glob2local_.clear();} 38 | void InitNode(const torch::Tensor& srcs, std::vector& out_nodes); 39 | void Induce(const std::vector& nodes, 40 | bool with_edge, 41 | std::vector& out_rows, 42 | std::vector& out_cols, 43 | std::vector& out_eids); 44 | private: 45 | IntHashMap glob2local_; 46 | }; 47 | 48 | } // namespace graphlearn_torch 49 | 50 | #endif // GRAPHLEARN_TORCH_CSRC_CPU_SUBGRAPH_OP_H_ 51 | -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cpu/weighted_sampler.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CPU_WEIGHTED_SAMPLER_H_ 17 | #define GRAPHLEARN_TORCH_CPU_WEIGHTED_SAMPLER_H_ 18 | 19 | #include "graphlearn_torch/include/sampler.h" 20 | 21 | 22 | namespace graphlearn_torch { 23 | 24 | class CPUWeightedSampler : public Sampler { 25 | public: 26 | CPUWeightedSampler(const Graph* graph) : Sampler(graph) {} 27 | ~CPUWeightedSampler() {} 28 | 29 | std::tuple Sample( 30 | const torch::Tensor& nodes, int32_t req_num) override; 31 | 32 | std::tuple 33 | SampleWithEdge(const torch::Tensor& nodes, int32_t req_num) override; 34 | 35 | private: 36 | void FillNbrsNum(const int64_t* nodes, const int32_t bs, 37 | const int32_t req_num, const int64_t row_count, 38 | const int64_t* row_ptr, int64_t* out_nbr_num); 39 | 40 | void CSRRowWiseSample(const int64_t* nodes, const int64_t* nbrs_offset, 41 | const int32_t bs, const int32_t req_num, const int64_t row_count, 42 | const int64_t* row_ptr, const int64_t* col_idx, const float* prob, 43 | int64_t* out_nbrs); 44 | 45 | void CSRRowWiseSample(const int64_t* nodes, const int64_t* nbrs_offset, 46 | const int32_t bs, const int32_t req_num, const int64_t row_count, 47 | const int64_t* row_ptr, const int64_t* col_idx, const float* prob, 48 | const int64_t* edge_id, int64_t* out_nbrs, int64_t* out_eid); 49 | 50 | 51 | void WeightedSample(const int64_t* col_begin, const int64_t* col_end, 52 | const int32_t req_num, const float* prob_begin, const float* prob_end, 53 | int64_t* out_nbrs); 54 | 55 | void WeightedSample(const int64_t* col_begin, const int64_t* col_end, 56 | const int64_t* eid_begin, const int64_t* eid_end, 57 | const int32_t req_num, const float* prob_begin, const float* prob_end, 58 | int64_t* out_nbrs, int64_t* out_eid); 59 | }; 60 | 61 | } // namespace graphlearn_torch 62 | 63 | #endif // GRAPHLEARN_TORCH_CPU_WEIGHTED_SAMPLER_H_ -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cuda/random_negative_sampler.cuh: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CUDA_RANDOM_NEGATIVE_SAMPLER_CUH_ 17 | #define GRAPHLEARN_TORCH_CUDA_RANDOM_NEGATIVE_SAMPLER_CUH_ 18 | 19 | #include "graphlearn_torch/include/negative_sampler.h" 20 | 21 | 22 | namespace graphlearn_torch { 23 | 24 | class CUDARandomNegativeSampler : public NegativeSampler { 25 | public: 26 | CUDARandomNegativeSampler(const Graph* graph) : NegativeSampler(graph) {} 27 | ~CUDARandomNegativeSampler() {} 28 | 29 | virtual std::tuple Sample( 30 | int32_t req_num, int32_t trials_num, bool padding=false); 31 | }; 32 | 33 | } // namespace graphlearn_torch 34 | 35 | #endif // GRAPHLEARN_TORCH_CUDA_RANDOM_NEGATIVE_SAMPLER_CUH_ -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cuda/random_sampler.cuh: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CUDA_RANDOM_SAMPLER_CUH_ 17 | #define GRAPHLEARN_TORCH_CUDA_RANDOM_SAMPLER_CUH_ 18 | 19 | #include "graphlearn_torch/include/sampler.h" 20 | 21 | namespace graphlearn_torch { 22 | 23 | class CUDARandomSampler : public Sampler { 24 | public: 25 | CUDARandomSampler(const Graph* graph) : Sampler(graph) {} 26 | ~CUDARandomSampler() {} 27 | 28 | std::tuple 29 | Sample(const torch::Tensor& nodes, int32_t req_num) override; 30 | 31 | std::tuple 32 | SampleWithEdge(const torch::Tensor& nodes, int32_t req_num) override; 33 | 34 | void CalNbrProb(int k, const torch::Tensor& last_prob, 35 | const torch::Tensor& nbr_last_prob, 36 | const Graph* nbr_graph_, torch::Tensor cur_prob); 37 | }; 38 | 39 | } // namespace graphlearn_torch 40 | 41 | #endif // GRAPHLEARN_TORCH_CUDA_RANDOM_SAMPLER_CUH_ -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cuda/subgraph_op.cuh: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CSRC_CUDA_SUBGRAPH_OP_CUH_ 17 | #define GRAPHLEARN_TORCH_CSRC_CUDA_SUBGRAPH_OP_CUH_ 18 | 19 | #include "graphlearn_torch/include/subgraph_op_base.h" 20 | 21 | #include 22 | #include 23 | #include 24 | 25 | #include "graphlearn_torch/include/graph.h" 26 | #include "graphlearn_torch/include/types.h" 27 | 28 | namespace graphlearn_torch { 29 | 30 | class HostHashTable; 31 | 32 | class CUDASubGraphOp : public SubGraphOp { 33 | public: 34 | CUDASubGraphOp(const Graph* graph); 35 | ~CUDASubGraphOp(); 36 | CUDASubGraphOp(const CUDASubGraphOp& other) = delete; 37 | CUDASubGraphOp& operator=(const CUDASubGraphOp& other) = delete; 38 | CUDASubGraphOp(CUDASubGraphOp&&) = default; 39 | CUDASubGraphOp& operator=(CUDASubGraphOp&&) = default; 40 | 41 | SubGraph NodeSubGraph(const torch::Tensor& srcs, bool with_edge=false) override; 42 | 43 | private: 44 | void Reset(); 45 | void InitNode(cudaStream_t stream, const torch::Tensor& srcs, 46 | int64_t* nodes, int32_t* nodes_size); 47 | void CSRSliceRows(cudaStream_t stream, 48 | const int64_t* rows, 49 | int32_t rows_size, 50 | int64_t* indptr); 51 | void GetNbrsNumAndColMask(cudaStream_t stream, 52 | const int64_t* nodes, 53 | const int64_t* origin_nbrs_offset, 54 | int32_t nodes_size, 55 | int64_t* out_nbrs_num); 56 | private: 57 | HostHashTable* host_table_; 58 | int32_t* col_mask_; 59 | }; 60 | 61 | } // namespace graphlearn_torch 62 | 63 | #endif // GRAPHLEARN_TORCH_CSRC_CUDA_SUBGRAPH_OP_CUH_ 64 | -------------------------------------------------------------------------------- /graphlearn_torch/csrc/cuda/weighted_sampler.cuh: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_CUDA_WEIGHTED_SAMPLER_CUH_ 17 | #define GRAPHLEARN_TORCH_CUDA_WEIGHTED_SAMPLER_CUH_ 18 | 19 | #include "graphlearn_torch/include/sampler.h" 20 | 21 | namespace graphlearn_torch { 22 | 23 | class CUDAWeightedSampler : public Sampler { 24 | public: 25 | CUDAWeightedSampler(const Graph* graph) : Sampler(graph) {} 26 | ~CUDAWeightedSampler() {} 27 | 28 | std::tuple Sample( 29 | const torch::Tensor& nodes, int32_t req_num) override { 30 | std::cerr << "Not supported yet!" << std::endl; 31 | } 32 | 33 | std::tuple 34 | SampleWithEdge(const torch::Tensor& nodes, int32_t req_num) override { 35 | std::cerr << "Not supported yet!" << std::endl; 36 | } 37 | 38 | }; 39 | 40 | } // namespace graphlearn_torch 41 | 42 | #endif // GRAPHLEARN_TORCH_CUDA_WEIGHTED_SAMPLER_CUH_ -------------------------------------------------------------------------------- /graphlearn_torch/csrc/sample_queue.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include "graphlearn_torch/include/sample_queue.h" 17 | 18 | namespace graphlearn_torch { 19 | 20 | void SampleQueue::Enqueue(const TensorMap& msg) { 21 | auto serialized_size = TensorMapSerializer::GetSerializedSize(msg); 22 | shmq_->Enqueue(serialized_size, [&msg] (void* buf) { 23 | TensorMapSerializer::Serialize(msg, buf); 24 | }); 25 | } 26 | 27 | TensorMap SampleQueue::Dequeue(unsigned int timeout_ms) { 28 | auto shm_data = shmq_->Dequeue(timeout_ms); 29 | return TensorMapSerializer::Load(std::move(shm_data)); 30 | } 31 | 32 | bool SampleQueue::Empty() { 33 | return shmq_->Empty(); 34 | } 35 | 36 | } // namespace graphlearn_torch 37 | -------------------------------------------------------------------------------- /graphlearn_torch/include/common.cuh: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_COMMON_CUH_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_COMMON_CUH_ 18 | 19 | #include 20 | #include 21 | #include 22 | #include 23 | #include 24 | 25 | #include "graphlearn_torch/include/common.h" 26 | 27 | 28 | namespace graphlearn_torch 29 | { 30 | 31 | inline void CUDARegisterByBlock(void* ptr, size_t size, unsigned int flags) { 32 | constexpr size_t BLOCK = 1000000000; 33 | auto* register_ptr = static_cast(ptr); 34 | for (size_t pos = 0; pos < size; pos += BLOCK) { 35 | auto s = std::min(size - pos, BLOCK); 36 | cudaHostRegister(register_ptr + pos, s, flags); 37 | } 38 | } 39 | 40 | inline void* CUDAAlloc(size_t nbytes, cudaStream_t stream = 0) { 41 | at::globalContext().lazyInitCUDA(); 42 | return c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(nbytes, stream); 43 | } 44 | 45 | inline void CUDADelete(void* ptr) { 46 | c10::cuda::CUDACachingAllocator::raw_delete(ptr); 47 | } 48 | 49 | #define CUDACheckError() \ 50 | { \ 51 | cudaError_t e = cudaGetLastError(); \ 52 | if ((e != cudaSuccess) && (e != cudaErrorPeerAccessAlreadyEnabled)) { \ 53 | printf("CUDA failure %s:%d: '%s'\n", __FILE__, __LINE__, \ 54 | cudaGetErrorString(e)); \ 55 | exit(EXIT_FAILURE); \ 56 | } \ 57 | } 58 | 59 | class CUDAAllocator { 60 | public: 61 | typedef char value_type; 62 | 63 | explicit CUDAAllocator(cudaStream_t stream) : stream_(stream) {} 64 | 65 | void operator()(void* ptr) const { 66 | CUDADelete(ptr); 67 | } 68 | 69 | template 70 | std::unique_ptr alloc_unique(std::size_t size) const { 71 | return std::unique_ptr( 72 | reinterpret_cast(CUDAAlloc(size, stream_)), *this); 73 | } 74 | 75 | char* allocate(std::ptrdiff_t size) const { 76 | return reinterpret_cast(CUDAAlloc(size, stream_)); 77 | } 78 | 79 | void deallocate(char* ptr, std::size_t) const { 80 | CUDADelete(ptr); 81 | } 82 | private: 83 | cudaStream_t stream_; 84 | }; 85 | 86 | } // namespace graphlearn_torch 87 | 88 | #endif // GRAPHLEARN_TORCH_INCLUDE_COMMON_CUH_ 89 | -------------------------------------------------------------------------------- /graphlearn_torch/include/common.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_COMMON_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_COMMON_H_ 18 | 19 | #include 20 | #include 21 | 22 | namespace graphlearn_torch 23 | { 24 | 25 | inline void Check(bool val, const char* err_msg) { 26 | if (val) { return; } 27 | throw std::runtime_error(err_msg); 28 | } 29 | 30 | template 31 | inline void CheckEq(const T &x, const T &y) { 32 | if (x == y) { return; } 33 | throw std::runtime_error(std::string("CheckEq failed")); 34 | } 35 | 36 | class RandomSeedManager { 37 | public: 38 | static RandomSeedManager& getInstance() { 39 | static RandomSeedManager instance; 40 | return instance; 41 | } 42 | 43 | void setSeed(uint32_t seed) { 44 | this->is_set = true; 45 | this->seed = seed; 46 | } 47 | 48 | uint32_t getSeed() const { 49 | if (this->is_set) { 50 | return seed; 51 | } 52 | else { 53 | std::random_device rd; 54 | return rd(); 55 | } 56 | } 57 | 58 | private: 59 | RandomSeedManager() {} // Constructor is private 60 | RandomSeedManager(RandomSeedManager const&) = delete; // Prevent copies 61 | void operator=(RandomSeedManager const&) = delete; // Prevent assignments 62 | 63 | uint32_t seed; 64 | bool is_set = false; 65 | }; 66 | 67 | 68 | } // namespace graphlearn_torch 69 | 70 | #endif // GRAPHLEARN_TORCH_INCLUDE_COMMON_H_ 71 | -------------------------------------------------------------------------------- /graphlearn_torch/include/func_factory.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_FUNC_FACTORY_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_FUNC_FACTORY_H_ 18 | 19 | #include 20 | #include 21 | 22 | namespace graphlearn_torch { 23 | 24 | /// A factory for function call with the same return type and argument types. 25 | template 26 | class FunctionFactory { 27 | public: 28 | typedef std::function Function; 29 | typedef std::unordered_map FactoryMap; 30 | 31 | static FunctionFactory& Get() { 32 | static FunctionFactory instance; 33 | return instance; 34 | } 35 | 36 | FunctionFactory(const FunctionFactory&) = delete; 37 | FunctionFactory& operator=(const FunctionFactory&) = delete; 38 | FunctionFactory(FunctionFactory&&) = delete; 39 | FunctionFactory& operator=(FunctionFactory&&) = delete; 40 | 41 | bool Register(const KeyType& key, Function func) { 42 | if (factory_map_.find(key) != factory_map_.end()) { 43 | return false; 44 | } 45 | factory_map_[key] = std::move(func); 46 | return true; 47 | } 48 | 49 | template < 50 | typename _RT = FuncReturnType, 51 | std::enable_if_t::value, int> = 0 52 | > 53 | auto Dispatch(const KeyType& key, FuncArgs... func_args) { 54 | if (factory_map_.find(key) == factory_map_.end()) { 55 | return _RT{}; 56 | } 57 | return factory_map_[key](std::forward(func_args)...); 58 | } 59 | 60 | template < 61 | typename _RT = FuncReturnType, 62 | std::enable_if_t::value, int> = 0 63 | > 64 | void Dispatch(const KeyType& key, FuncArgs... func_args) { 65 | if (factory_map_.find(key) == factory_map_.end()) { 66 | return; 67 | } 68 | factory_map_[key](std::forward(func_args)...); 69 | } 70 | 71 | private: 72 | FactoryMap factory_map_; 73 | 74 | FunctionFactory() = default; 75 | }; 76 | 77 | } // namespace graphlearn_torch 78 | 79 | #endif // GRAPHLEARN_TORCH_INCLUDE_FUNC_FACTORY_H_ 80 | -------------------------------------------------------------------------------- /graphlearn_torch/include/inducer_base.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_INDUCER_BASE_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_INDUCER_BASE_H_ 18 | 19 | #include 20 | #include 21 | #include 22 | 23 | #include "graphlearn_torch/include/types.h" 24 | 25 | namespace graphlearn_torch { 26 | 27 | 28 | class Inducer { 29 | public: 30 | Inducer() {} 31 | virtual ~Inducer() {} 32 | 33 | // Init inducer with seed nodes and return de-duplicated seed nodes. 34 | virtual torch::Tensor InitNode(const torch::Tensor& seed) = 0; 35 | 36 | // Induce COO subgraph incrementally. 37 | // Returns: incremental unique nodes, rows, cols. 38 | // Note that the current srcs must be a subset of last output nodes. 39 | virtual std::tuple 40 | InduceNext(const torch::Tensor& srcs, 41 | const torch::Tensor& nbrs, 42 | const torch::Tensor& nbrs_num) = 0; 43 | 44 | virtual void Reset() = 0; 45 | }; 46 | 47 | 48 | class HeteroInducer { 49 | public: 50 | HeteroInducer() {} 51 | virtual ~HeteroInducer() {} 52 | 53 | // Init inducer with seed nodes and return de-duplicated seed nodes. 54 | virtual TensorMap InitNode(const TensorMap& seed) = 0; 55 | 56 | // Induce HeteroCOO subgraph incrementally. 57 | // Returns: incremental unique nodes_dict, rows_dict, cols_dict. 58 | // Note that each type of the current srcs must be a subset of last 59 | // corresponding type output nodes. 60 | virtual HeteroCOO InduceNext(const HeteroNbr& nbrs) = 0; 61 | 62 | virtual void Reset() = 0; 63 | }; 64 | 65 | } // namespace graphlearn_torch 66 | 67 | #endif // GRAPHLEARN_TORCH_INCLUDE_INDUCER_BASE_H_ 68 | -------------------------------------------------------------------------------- /graphlearn_torch/include/negative_sampler.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_NEGATIVE_SAMPLER_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_NEGATIVE_SAMPLER_H_ 18 | 19 | #include "graphlearn_torch/include/graph.h" 20 | 21 | namespace graphlearn_torch { 22 | 23 | class NegativeSampler { 24 | public: 25 | NegativeSampler(const Graph* graph) : graph_(graph) {} 26 | virtual ~NegativeSampler() {} 27 | 28 | virtual std::tuple Sample( 29 | int32_t req_num, int32_t trials_num, bool padding=false) = 0; 30 | 31 | protected: 32 | const Graph* graph_; 33 | }; 34 | 35 | } // namespace graphlearn_torch 36 | 37 | #endif // GRAPHLEARN_TORCH_INCLUDE_NEGATIVE_SAMPLER_H_ -------------------------------------------------------------------------------- /graphlearn_torch/include/sample_queue.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_SAMPLE_QUEUE_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_SAMPLE_QUEUE_H_ 18 | 19 | #include 20 | 21 | #include "graphlearn_torch/include/shm_queue.h" 22 | #include "graphlearn_torch/include/tensor_map.h" 23 | 24 | namespace graphlearn_torch { 25 | 26 | class SampleQueue { 27 | public: 28 | SampleQueue(size_t max_msg_num, size_t buf_size) { 29 | shmq_ = std::make_unique(max_msg_num, buf_size); 30 | } 31 | 32 | SampleQueue(int shmid) { 33 | shmq_ = std::make_unique(shmid); 34 | } 35 | 36 | void PinMemory() { 37 | shmq_->PinMemory(); 38 | } 39 | 40 | int ShmId() const { 41 | return shmq_->ShmId(); 42 | } 43 | 44 | void Enqueue(const TensorMap& msg); 45 | TensorMap Dequeue(unsigned int timeout_ms = 0); 46 | bool Empty(); 47 | 48 | private: 49 | std::unique_ptr shmq_; 50 | }; 51 | 52 | } // namespace graphlearn_torch 53 | 54 | #endif // GRAPHLEARN_TORCH_INCLUDE_SAMPLE_QUEUE_H_ 55 | -------------------------------------------------------------------------------- /graphlearn_torch/include/sampler.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_SAMPLER_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_SAMPLER_H_ 18 | 19 | #include "graphlearn_torch/include/graph.h" 20 | 21 | namespace graphlearn_torch { 22 | 23 | class Sampler { 24 | public: 25 | Sampler(const Graph* graph) : graph_(graph) {} 26 | virtual ~Sampler() {} 27 | 28 | // return: (nbrs, nbrs_num) 29 | virtual std::tuple Sample( 30 | const torch::Tensor& nodes, int32_t req_num) = 0; 31 | 32 | // return: (nbrs, nbrs_num, edge_ids) 33 | virtual std::tuple 34 | SampleWithEdge(const torch::Tensor& nodes, int32_t req_num) = 0; 35 | 36 | protected: 37 | const Graph* graph_; 38 | }; 39 | 40 | } // namespace graphlearn_torch 41 | 42 | #endif // GRAPHLEARN_TORCH_INCLUDE_SAMPLER_H_ -------------------------------------------------------------------------------- /graphlearn_torch/include/stitch_sample_results.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_STITCH_SAMPLE_RESULTS_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_STITCH_SAMPLE_RESULTS_H_ 18 | 19 | #include 20 | 21 | namespace graphlearn_torch { 22 | 23 | std::tuple> 24 | CPUStitchSampleResults(const torch::Tensor& ids, 25 | const std::vector& idx_list, 26 | const std::vector& nbrs_list, 27 | const std::vector& nbrs_num_list, 28 | const std::vector& eids_list); 29 | 30 | std::tuple> 31 | CUDAStitchSampleResults(const torch::Tensor& ids, 32 | const std::vector& idx_list, 33 | const std::vector& nbrs_list, 34 | const std::vector& nbrs_num_list, 35 | const std::vector& eids_list); 36 | 37 | } // namespace graphlearn_torch 38 | 39 | #endif // GRAPHLEARN_TORCH_INCLUDE_STITCH_SAMPLE_RESULTS_H_ 40 | -------------------------------------------------------------------------------- /graphlearn_torch/include/subgraph_op_base.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_SUBGRAPH_OP_BASE_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_SUBGRAPH_OP_BASE_H_ 18 | 19 | #include "graphlearn_torch/include/graph.h" 20 | #include "graphlearn_torch/include/types.h" 21 | 22 | namespace graphlearn_torch { 23 | 24 | class SubGraphOp { 25 | public: 26 | SubGraphOp(const Graph* graph) : graph_(graph) {} 27 | virtual ~SubGraphOp() {} 28 | 29 | virtual SubGraph NodeSubGraph(const torch::Tensor& srcs, 30 | bool with_edge=false) = 0; 31 | protected: 32 | const Graph* graph_; 33 | }; 34 | 35 | } // namespace graphlearn_torch 36 | 37 | #endif // GRAPHLEARN_TORCH_INCLUDE_SUBGRAPH_OP_BASE_H_ -------------------------------------------------------------------------------- /graphlearn_torch/include/tensor_map.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_TENSOR_MAP_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_TENSOR_MAP_H_ 18 | 19 | #include "graphlearn_torch/include/shm_queue.h" 20 | #include "graphlearn_torch/include/types.h" 21 | 22 | namespace graphlearn_torch { 23 | 24 | /// Serialization layout of a `TensorMap`: 25 | /// | tensor_num | serialized_tensor_1 | ... | serialized_tensor_n | 26 | /// 27 | /// Serialization layout of a `torch::Tensor`: 28 | /// | key_length | key_bytes | data_type | shape_num | shape_values | data_length | data_bytes | 29 | 30 | class TensorMapSerializer { 31 | public: 32 | /// Get the serialized size of a `TensorMap` with all its managed tensors. 33 | /// 34 | static size_t GetSerializedSize(const TensorMap& map); 35 | 36 | /// Serialize a `TensorMap` into a data buffer. 37 | /// 38 | static void Serialize(const TensorMap& map, void* buf); 39 | 40 | /// Create a `TensorMap` from a serialized data buffer without taking 41 | /// ownership of the original data. 42 | /// 43 | /// The original data should maintain its lifetime during the use of 44 | /// created `TensorMap`. 45 | /// 46 | static TensorMap Load(void* buf); 47 | 48 | /// Create a `TensorMap` from a `ShmData`, the created `TensorMap` will 49 | /// take ownership of this shm data and release it after use. 50 | /// 51 | static TensorMap Load(ShmData&& shm_data); 52 | }; 53 | 54 | } // namespace graphlearn_torch 55 | 56 | #endif // GRAPHLEARN_TORCH_INCLUDE_TENSOR_MAP_H_ 57 | -------------------------------------------------------------------------------- /graphlearn_torch/include/types.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_TYPES_H_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_TYPES_H_ 18 | 19 | #include 20 | #include 21 | #include 22 | #include 23 | #include 24 | 25 | namespace graphlearn_torch 26 | { 27 | 28 | using TensorMap = std::unordered_map; 29 | using IntHashMap = std::unordered_map; 30 | using EdgeType = std::tuple; 31 | 32 | struct EdgeTypeHash 33 | { 34 | std::size_t operator()(const EdgeType& t) const noexcept 35 | { 36 | auto edge_type = std::get<0>(t) + "_" + std::get<1>(t) +\ 37 | "_" + std::get<2>(t); 38 | return std::hash{}(edge_type); 39 | } 40 | }; 41 | 42 | using TensorEdgeMap = std::unordered_map; 43 | // node_dict, row_dict, col_dict. 44 | using HeteroCOO = std::tuple; 45 | // dict of src, nbrs, nbrs_num. 46 | using HeteroNbr = std::unordered_map, EdgeTypeHash>; 48 | 49 | template 50 | struct Array 51 | { 52 | Array() {} 53 | Array(const T* ptr, const int64_t ptr_size) { 54 | data = ptr; 55 | size = ptr_size; 56 | } 57 | const T* data; 58 | int64_t size; 59 | }; 60 | 61 | struct SubGraph { 62 | SubGraph() {} 63 | SubGraph(torch::Tensor nodes, torch::Tensor rows, torch::Tensor cols): 64 | nodes(nodes), rows(rows), cols(cols) {} 65 | 66 | torch::Tensor nodes; 67 | torch::Tensor rows; 68 | torch::Tensor cols; 69 | torch::Tensor eids; 70 | }; 71 | 72 | } // namespace graphlearn_torch 73 | 74 | #endif // GRAPHLEARN_TORCH_INCLUDE_TYPES_H_ -------------------------------------------------------------------------------- /graphlearn_torch/include/unified_tensor.cuh: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #ifndef GRAPHLEARN_TORCH_INCLUDE_UNIFIED_TENSOR_CUH_ 17 | #define GRAPHLEARN_TORCH_INCLUDE_UNIFIED_TENSOR_CUH_ 18 | 19 | #include 20 | #include 21 | #include 22 | #include 23 | 24 | #include 25 | #include 26 | #include 27 | 28 | namespace graphlearn_torch { 29 | 30 | // CUDA Ipc shared tensor. 31 | class SharedTensor { 32 | public: 33 | SharedTensor(); 34 | SharedTensor(int32_t device, cudaIpcMemHandle_t mem_handle, 35 | const std::vector& shape); 36 | std::tuple> ShareCUDAIpc(); 37 | void FromCUDAIpc( 38 | std::tuple> ipc_data); 39 | 40 | public: 41 | int32_t device_; 42 | cudaIpcMemHandle_t mem_handle_; 43 | std::vector shape_; 44 | }; 45 | 46 | class UnifiedTensor { 47 | public: 48 | UnifiedTensor(int32_t device, torch::Dtype dtype = torch::kFloat32); 49 | ~UnifiedTensor(); 50 | void AppendSharedTensor(const SharedTensor& item); 51 | void AppendCPUTensor(const torch::Tensor& tensor); 52 | void InitFrom(const std::vector& tensors, 53 | const std::vector& tensor_devices); 54 | torch::Tensor operator[](const torch::Tensor& indices); 55 | std::vector ShareCUDAIpc(); 56 | 57 | const std::vector Shape() const; 58 | int32_t Device() const; 59 | int32_t Size(int32_t dim) const; 60 | int64_t Stride(int32_t dim) const; 61 | int64_t Numel() const; 62 | 63 | std::tuple GetDeviceData(int32_t device); 64 | void Initp2p(const std::vector& devices); 65 | 66 | private: 67 | std::vector shape_; 68 | torch::Dtype dtype_; 69 | int32_t device_; 70 | int32_t device_count_; 71 | bool inited_; 72 | void* registered_ptr_; 73 | 74 | std::vector tensor_offsets_; 75 | std::vector> tensor_shapes_; 76 | std::vector tensor_ptrs_; 77 | std::vector tensor_devices_; 78 | std::unordered_map> device_data_map_; 79 | }; 80 | 81 | } // namespace graphlearn_torch 82 | 83 | #endif // GRAPHLEARN_TORCH_INCLUDE_UNIFIED_TENSOR_CUH_ 84 | -------------------------------------------------------------------------------- /graphlearn_torch/python/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import os 17 | 18 | # Disable warning messages of LibTorch by default. 19 | os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR" 20 | 21 | from . import channel 22 | from . import data 23 | from . import distributed 24 | from . import loader 25 | from . import partition 26 | from . import sampler 27 | from . import utils 28 | from .typing import * 29 | -------------------------------------------------------------------------------- /graphlearn_torch/python/channel/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from .base import SampleMessage, ChannelBase, QueueTimeoutError 17 | from .mp_channel import MpChannel 18 | from .shm_channel import ShmChannel 19 | from .remote_channel import RemoteReceivingChannel 20 | -------------------------------------------------------------------------------- /graphlearn_torch/python/channel/base.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from abc import ABC, abstractmethod 17 | from typing import Dict 18 | import torch 19 | from .. import py_graphlearn_torch as pywrap 20 | 21 | QueueTimeoutError = pywrap.QueueTimeoutError 22 | 23 | # A `SampleMessage` contains all possible results from a sampler, including 24 | # subgraph data, features and user defined metas. 25 | SampleMessage = Dict[str, torch.Tensor] 26 | 27 | 28 | class ChannelBase(ABC): 29 | r""" A base class that initializes a channel for sample messages and 30 | provides :meth:`send` and :meth:`recv` routines. 31 | """ 32 | @abstractmethod 33 | def send(self, msg: SampleMessage, **kwargs): 34 | r""" Send a sample message into channel, the implemented channel should 35 | porcess this message data properly. 36 | 37 | Args: 38 | msg: The sample message to send. 39 | """ 40 | 41 | @abstractmethod 42 | def recv(self, **kwargs) -> SampleMessage: 43 | r""" Recv a sample message from channel. 44 | """ 45 | -------------------------------------------------------------------------------- /graphlearn_torch/python/channel/mp_channel.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import torch.multiprocessing as mp 17 | 18 | from .base import SampleMessage, ChannelBase 19 | 20 | 21 | class MpChannel(ChannelBase): 22 | r""" A simple multiprocessing channel using `torch.multiprocessing.Queue`. 23 | 24 | Args: 25 | The input arguments should be consistent with `torch.multiprocessing.Queue`. 26 | """ 27 | def __init__(self, **kwargs): 28 | self._queue = mp.get_context('spawn').Queue(**kwargs) 29 | 30 | def send(self, msg: SampleMessage, **kwargs): 31 | self._queue.put(msg, **kwargs) 32 | 33 | def recv(self, **kwargs) -> SampleMessage: 34 | return self._queue.get(**kwargs) 35 | -------------------------------------------------------------------------------- /graphlearn_torch/python/channel/shm_channel.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from typing import Union 17 | 18 | from .. import py_graphlearn_torch as pywrap 19 | from ..utils import parse_size 20 | 21 | from .base import SampleMessage, ChannelBase 22 | 23 | 24 | class ShmChannel(ChannelBase): 25 | r""" A communication channel for sample messages based on a shared-memory 26 | queue, which is implemented in the underlying c++ lib. 27 | 28 | Note that the underlying shared-memory buffer of this channel is pinnable, 29 | which will achieve better performance when the consumer needs to copy 30 | data from channel to gpu device. 31 | 32 | Args: 33 | capacity: The max bufferd number of sample messages in channel. 34 | shm_size: The allocated size (bytes) for underlying shared-memory. 35 | 36 | When the producer send sample message to the channel, it will be limited by 37 | both `capacity` and `shm_size`. E.g, if current number of buffered 38 | messages in channel reaches the `capacity` limit, or current used 39 | buffer memory reaches the `shm_size` limit, the current `send` operation 40 | will be blocked until some messages in channel are consumed and related 41 | resource are released. 42 | """ 43 | def __init__(self, 44 | capacity: int=128, 45 | shm_size: Union[str, int]='256MB'): 46 | assert capacity > 0 47 | shm_size = parse_size(shm_size) 48 | self._queue = pywrap.SampleQueue(capacity, shm_size) 49 | 50 | def pin_memory(self): 51 | r""" Pin underlying shared-memory. 52 | """ 53 | self._queue.pin_memory() 54 | 55 | def empty(self) -> bool: 56 | r""" Whether the queue is empty. 57 | """ 58 | return self._queue.empty() 59 | 60 | def send(self, msg: SampleMessage, **kwargs): 61 | self._queue.send(msg) 62 | 63 | def recv(self, timeout_ms=None, **kwargs) -> SampleMessage: 64 | if timeout_ms is None: 65 | timeout_ms = 0 66 | return self._queue.receive(timeout_ms=timeout_ms) 67 | -------------------------------------------------------------------------------- /graphlearn_torch/python/data/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from .dataset import * 17 | from .feature import * 18 | from .graph import * 19 | from .table_dataset import * 20 | from .reorder import * 21 | from .unified_tensor import * 22 | from .vineyard_utils import * 23 | -------------------------------------------------------------------------------- /graphlearn_torch/python/data/reorder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import torch 17 | 18 | 19 | def sort_by_in_degree(cpu_tensor, shuffle_ratio, topo): 20 | if topo is None: 21 | return cpu_tensor, None 22 | row_count = topo.row_count 23 | assert not (cpu_tensor.size(0) < row_count) 24 | new_idx = torch.arange(row_count, dtype=torch.long) 25 | perm_range = torch.randperm(int(row_count * shuffle_ratio)) 26 | _, old_idx = torch.sort(topo.degrees, descending=True) 27 | old2new = torch.arange(cpu_tensor.size(0), dtype=torch.long) 28 | old_idx[:int(row_count * shuffle_ratio)] = old_idx[perm_range] 29 | tmp_t = cpu_tensor[old_idx] 30 | if row_count < cpu_tensor.size(0): 31 | cpu_tensor = torch.cat([tmp_t, cpu_tensor[row_count:]], dim=0) 32 | else: 33 | cpu_tensor = tmp_t 34 | old2new[old_idx] = new_idx 35 | del new_idx, perm_range, old_idx, tmp_t 36 | return cpu_tensor, old2new 37 | -------------------------------------------------------------------------------- /graphlearn_torch/python/distributed/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from .dist_client import ( 17 | init_client, shutdown_client, async_request_server, request_server 18 | ) 19 | from .dist_context import DistRole, DistContext, get_context, init_worker_group 20 | from .dist_dataset import DistDataset 21 | from .dist_feature import PartialFeature, DistFeature 22 | from .dist_graph import DistGraph 23 | from .dist_link_neighbor_loader import DistLinkNeighborLoader 24 | from .dist_loader import DistLoader 25 | from .dist_neighbor_loader import DistNeighborLoader 26 | from .dist_neighbor_sampler import DistNeighborSampler 27 | from .dist_options import ( 28 | CollocatedDistSamplingWorkerOptions, 29 | MpDistSamplingWorkerOptions, 30 | RemoteDistSamplingWorkerOptions 31 | ) 32 | from .dist_random_partitioner import DistRandomPartitioner 33 | from .dist_sampling_producer import ( 34 | DistMpSamplingProducer, DistCollocatedSamplingProducer 35 | ) 36 | from .dist_server import ( 37 | DistServer, get_server, init_server, wait_and_shutdown_server 38 | ) 39 | from .dist_subgraph_loader import DistSubGraphLoader 40 | from .dist_table_dataset import DistTableDataset, DistTableRandomPartitioner 41 | from .event_loop import ConcurrentEventLoop 42 | from .rpc import ( 43 | init_rpc, shutdown_rpc, rpc_is_initialized, 44 | get_rpc_master_addr, get_rpc_master_port, 45 | all_gather, barrier, global_all_gather, global_barrier, 46 | RpcDataPartitionRouter, rpc_sync_data_partitions, 47 | RpcCalleeBase, rpc_register, rpc_request_async, rpc_request, 48 | rpc_global_request_async, rpc_global_request 49 | ) 50 | -------------------------------------------------------------------------------- /graphlearn_torch/python/distributed/event_loop.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import asyncio 17 | import logging 18 | from threading import Thread, BoundedSemaphore 19 | 20 | import torch 21 | 22 | 23 | def wrap_torch_future(f: torch.futures.Future) -> asyncio.futures.Future: 24 | r""" Convert a torch future to a standard asyncio future. 25 | """ 26 | loop = asyncio.get_event_loop() 27 | aio_future = loop.create_future() 28 | def on_done(*_): 29 | try: 30 | result = f.wait() 31 | except Exception as e: 32 | loop.call_soon_threadsafe(aio_future.set_exception, e) 33 | else: 34 | loop.call_soon_threadsafe(aio_future.set_result, result) 35 | f.add_done_callback(on_done) 36 | return aio_future 37 | 38 | 39 | class ConcurrentEventLoop(object): 40 | r""" Concurrent event loop context. 41 | 42 | Args: 43 | concurrency: max processing concurrency. 44 | """ 45 | def __init__(self, concurrency): 46 | self._concurrency = concurrency 47 | self._sem = BoundedSemaphore(concurrency) 48 | self._loop = asyncio.new_event_loop() 49 | self._runner_t = Thread(target=self._run_loop) 50 | self._runner_t.daemon = True 51 | 52 | def start_loop(self): 53 | if not self._runner_t.is_alive(): 54 | self._runner_t.start() 55 | 56 | def shutdown_loop(self): 57 | self.wait_all() 58 | if self._runner_t.is_alive(): 59 | self._loop.stop() 60 | self._runner_t.join(timeout=1) 61 | 62 | def wait_all(self): 63 | r""" Wait all pending tasks to be finished. 64 | """ 65 | for _ in range(self._concurrency): 66 | self._sem.acquire() 67 | for _ in range(self._concurrency): 68 | self._sem.release() 69 | 70 | def add_task(self, coro, callback=None): 71 | r""" Add an asynchronized coroutine task to run. 72 | 73 | Args: 74 | coro: the async coroutine func. 75 | callback: the callback func applied on the returned results 76 | after the coroutine task is finished. 77 | 78 | Note that any results returned by `callback` func will be ignored, 79 | so it is preferable to handle all in your `callback` func and do 80 | not return any results. 81 | """ 82 | def on_done(f: asyncio.futures.Future): 83 | try: 84 | res = f.result() 85 | if callback is not None: 86 | callback(res) 87 | except Exception as e: 88 | logging.error("coroutine task failed: %s", e) 89 | self._sem.release() 90 | self._sem.acquire() 91 | fut = asyncio.run_coroutine_threadsafe(coro, self._loop) 92 | fut.add_done_callback(on_done) 93 | 94 | def run_task(self, coro): 95 | r""" Run a coroutine task synchronously. 96 | """ 97 | with self._sem: 98 | fut = asyncio.run_coroutine_threadsafe(coro, self._loop) 99 | return fut.result() 100 | 101 | def _run_loop(self): 102 | self._loop.run_forever() 103 | -------------------------------------------------------------------------------- /graphlearn_torch/python/loader/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from .link_loader import * 17 | from .link_neighbor_loader import * 18 | from .node_loader import * 19 | from .neighbor_loader import * 20 | from .transform import * 21 | -------------------------------------------------------------------------------- /graphlearn_torch/python/partition/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from .base import * 17 | from .frequency_partitioner import FrequencyPartitioner 18 | from .partition_book import * 19 | from .random_partitioner import RandomPartitioner 20 | -------------------------------------------------------------------------------- /graphlearn_torch/python/partition/partition_book.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import List, Tuple 3 | from .base import PartitionBook 4 | 5 | 6 | class RangePartitionBook(PartitionBook): 7 | r"""A class for managing range-based partitions of consecutive IDs. 8 | Suitable when IDs within each partition are consecutive. 9 | Args: 10 | partition_ranges (List[Tuple[int, int]]): A list of tuples representing 11 | the start and end (exclusive) of each partition range. 12 | partition_idx (int): The index of the current partition. 13 | Example: 14 | >>> partition_ranges = [(0, 10), (10, 20), (20, 30)] 15 | >>> range_pb = RangePartitionBook(partition_ranges, partition_idx=1) 16 | >>> indices = torch.tensor([0, 5, 10, 15, 20, 25]) 17 | >>> partition_ids = range_pb[indices] 18 | >>> print(partition_ids) 19 | tensor([0, 0, 1, 1, 2, 2]) 20 | """ 21 | 22 | def __init__(self, partition_ranges: List[Tuple[int, int]], partition_idx: int): 23 | if not all(r[0] < r[1] for r in partition_ranges): 24 | raise ValueError("All partition ranges must have start < end") 25 | if not all(r1[1] == r2[0] for r1, r2 in zip(partition_ranges[:-1], partition_ranges[1:])): 26 | raise ValueError("Partition ranges must be continuous") 27 | 28 | self.partition_bounds = torch.tensor( 29 | [end for _, end in partition_ranges], dtype=torch.long) 30 | self.partition_idx = partition_idx 31 | self._id2index = OffsetId2Index(partition_ranges[partition_idx][0]) 32 | 33 | def __getitem__(self, indices: torch.Tensor) -> torch.Tensor: 34 | return torch.searchsorted(self.partition_bounds, indices, right=True) 35 | 36 | @property 37 | def device(self): 38 | return self.partition_bounds.device 39 | 40 | @property 41 | def id2index(self): 42 | return self._id2index 43 | 44 | def id_filter(self, node_pb: PartitionBook, partition_idx: int): 45 | start = self.partition_bounds[partition_idx-1] if partition_idx > 0 else 0 46 | end = self.partition_bounds[partition_idx] 47 | return torch.arange(start, end) 48 | 49 | 50 | class OffsetId2Index: 51 | r""" 52 | Convert global IDs to local indices by subtracting a specified offset. 53 | """ 54 | 55 | def __init__(self, offset: int): 56 | self.offset = offset 57 | 58 | def __getitem__(self, ids: torch.Tensor) -> torch.Tensor: 59 | local_indices = ids - self.offset 60 | return local_indices 61 | 62 | def to(self, device): 63 | # device is always same as the input ids 64 | return self 65 | 66 | 67 | class GLTPartitionBook(PartitionBook, torch.Tensor): 68 | r""" A partition book of graph nodes or edges. 69 | """ 70 | 71 | def __getitem__(self, indices) -> torch.Tensor: 72 | return torch.Tensor.__getitem__(self, indices) 73 | -------------------------------------------------------------------------------- /graphlearn_torch/python/partition/random_partitioner.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | 17 | from typing import List, Dict, Optional, Tuple, Union 18 | 19 | import torch 20 | 21 | from ..typing import NodeType, EdgeType, TensorDataType 22 | 23 | from .base import PartitionerBase, PartitionBook 24 | 25 | 26 | # Implementation of a random partitioner. 27 | 28 | class RandomPartitioner(PartitionerBase): 29 | r""" Random partitioner for graph topology and features. 30 | 31 | Args: 32 | output_dir: The output root directory for partitioned results. 33 | num_parts: Number of partitions. 34 | num_nodes: Number of graph nodes, should be a dict for hetero data. 35 | edge_index: The edge index data of graph edges, should be a dict 36 | for hetero data. 37 | node_feat: The node feature data, should be a dict for hetero data. 38 | node_feat_dtype: The data type of node features. 39 | edge_feat: The edge feature data, should be a dict for hetero data. 40 | edge_feat_dtype: The data type of edge features. 41 | edge_assign_strategy: The assignment strategy when partitioning edges, 42 | should be 'by_src' or 'by_dst'. 43 | chunk_size: The chunk size for partitioning. 44 | """ 45 | def __init__( 46 | self, 47 | output_dir: str, 48 | num_parts: int, 49 | num_nodes: Union[int, Dict[NodeType, int]], 50 | edge_index: Union[TensorDataType, Dict[EdgeType, TensorDataType]], 51 | node_feat: Optional[Union[TensorDataType, Dict[NodeType, TensorDataType]]] = None, 52 | node_feat_dtype: torch.dtype = torch.float32, 53 | edge_feat: Optional[Union[TensorDataType, Dict[EdgeType, TensorDataType]]] = None, 54 | edge_feat_dtype: torch.dtype = torch.float32, 55 | edge_weights: Optional[Union[TensorDataType, Dict[EdgeType, TensorDataType]]] = None, 56 | edge_assign_strategy: str = 'by_src', 57 | chunk_size: int = 10000, 58 | ): 59 | super().__init__(output_dir, num_parts, num_nodes, edge_index, node_feat, 60 | node_feat_dtype, edge_feat, edge_feat_dtype, edge_weights, 61 | edge_assign_strategy, chunk_size) 62 | 63 | def _partition_node( 64 | self, 65 | ntype: Optional[NodeType] = None 66 | ) -> Tuple[List[torch.Tensor], PartitionBook]: 67 | if 'hetero' == self.data_cls: 68 | assert ntype is not None 69 | node_num = self.num_nodes[ntype] 70 | else: 71 | node_num = self.num_nodes 72 | ids = torch.arange(node_num, dtype=torch.int64) 73 | partition_book = ids % self.num_parts 74 | rand_order = torch.randperm(ids.size(0)) 75 | partition_book = partition_book[rand_order] 76 | partition_results = [] 77 | for pidx in range(self.num_parts): 78 | mask = (partition_book == pidx) 79 | partition_results.append(torch.masked_select(ids, mask)) 80 | return partition_results, partition_book 81 | 82 | def _cache_node( 83 | self, 84 | ntype: Optional[NodeType] = None 85 | ) -> List[Optional[torch.Tensor]]: 86 | return [None for _ in range(self.num_parts)] 87 | -------------------------------------------------------------------------------- /graphlearn_torch/python/py_export_v6d.cc: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include 17 | #include 18 | #include 19 | #include 20 | 21 | #include 22 | 23 | #include "graphlearn_torch/v6d/vineyard_utils.h" 24 | 25 | namespace py = pybind11; 26 | 27 | using namespace graphlearn_torch::vineyard_utils; 28 | 29 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 30 | m.doc() = "Python bindings for vineyard utils C++ frontend"; 31 | m.def("vineyard_to_csr", &ToCSR); 32 | m.def("load_vertex_feature_from_vineyard", &LoadVertexFeatures); 33 | m.def("load_edge_feature_from_vineyard", &LoadEdgeFeatures); 34 | m.def("get_frag_vertex_offset", &GetFragVertexOffset); 35 | m.def("get_frag_vertex_num", &GetFragVertexNum); 36 | 37 | py::class_(m, "VineyardFragHandle") 38 | .def(py::init()) 39 | .def("get_fid_from_gid", &VineyardFragHandle::GetFidFromGid, py::arg("gid")) 40 | .def("get_inner_vertices", &VineyardFragHandle::GetInnerVertices, 41 | py::arg("v_label_name")); 42 | } 43 | -------------------------------------------------------------------------------- /graphlearn_torch/python/sampler/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from .base import * 17 | from .negative_sampler import * 18 | from .neighbor_sampler import * -------------------------------------------------------------------------------- /graphlearn_torch/python/sampler/negative_sampler.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import torch 17 | 18 | from .. import py_graphlearn_torch as pywrap 19 | 20 | 21 | class RandomNegativeSampler(object): 22 | r""" Random negative Sampler. 23 | 24 | Args: 25 | graph: A ``graphlearn_torch.data.Graph`` object. 26 | mode: Execution mode of sampling, 'CUDA' means sampling on 27 | GPU, 'CPU' means sampling on CPU. 28 | edge_dir: The direction of edges to be sampled, determines 29 | the order of rows and columns returned. 30 | """ 31 | def __init__(self, graph, mode='CUDA', edge_dir='out'): 32 | self._mode = mode 33 | self.edge_dir = edge_dir 34 | if mode == 'CUDA': 35 | self._sampler = pywrap.CUDARandomNegativeSampler(graph.graph_handler) 36 | else: 37 | self._sampler = pywrap.CPURandomNegativeSampler(graph.graph_handler) 38 | 39 | def sample(self, req_num, trials_num=5, padding=False): 40 | r""" Negative sampling. 41 | 42 | Args: 43 | req_num: The number of request(max) negative samples. 44 | trials_num: The number of trials for negative sampling. 45 | padding: Whether to patch the negative sampling results to req_num. 46 | If True, after trying trials_num times, if the number of true negative 47 | samples is still less than req_num, just random sample edges(non-strict 48 | negative) as negative samples. 49 | 50 | Returns: 51 | negative edge_index(non-strict when padding is True). 52 | """ 53 | if self.edge_dir == 'out': 54 | rows, cols = self._sampler.sample(req_num, trials_num, padding) 55 | elif self.edge_dir == 'in': 56 | cols, rows = self._sampler.sample(req_num, trials_num, padding) 57 | return torch.stack([rows, cols], dim=0) 58 | -------------------------------------------------------------------------------- /graphlearn_torch/python/typing.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from typing import Dict, List, NamedTuple, Optional, Tuple, Union 17 | 18 | import torch 19 | import numpy as np 20 | from enum import Enum 21 | 22 | # Types for basic graph entity ################################################# 23 | 24 | # Node-types are denoted by a single string 25 | NodeType = str 26 | 27 | # Edge-types are denotes by a triplet of strings. 28 | EdgeType = Tuple[str, str, str] 29 | 30 | EDGE_TYPE_STR_SPLIT = '__' 31 | 32 | def as_str(type: Union[NodeType, EdgeType]) -> str: 33 | if isinstance(type, NodeType): 34 | return type 35 | elif isinstance(type, (list, tuple)) and len(type) == 3: 36 | return EDGE_TYPE_STR_SPLIT.join(type) 37 | return '' 38 | 39 | def reverse_edge_type(etype: EdgeType): 40 | src, edge, dst = etype 41 | if not src == dst: 42 | if edge.split("_", 1)[0] == 'rev': # undirected edge with `rev_` prefix. 43 | edge = edge.split("_", 1)[1] 44 | else: 45 | edge = 'rev_' + edge 46 | return (dst, edge, src) 47 | 48 | 49 | # A representation of tensor data 50 | TensorDataType = Union[torch.Tensor, np.ndarray] 51 | 52 | NodeLabel = Union[TensorDataType, Dict[NodeType, TensorDataType]] 53 | NodeIndex = Union[TensorDataType, Dict[NodeType, TensorDataType]] 54 | 55 | class Split(Enum): 56 | train = 'train' 57 | valid = 'valid' 58 | test = 'test' 59 | 60 | # Types for partition data ##################################################### 61 | 62 | class GraphPartitionData(NamedTuple): 63 | r""" Data and indexing info of a graph partition. 64 | """ 65 | # edge index (rows, cols) 66 | edge_index: Tuple[torch.Tensor, torch.Tensor] 67 | # edge ids tensor corresponding to `edge_index` 68 | eids: torch.Tensor 69 | # weights tensor corresponding to `edge_index` 70 | weights: Optional[torch.Tensor] = None 71 | 72 | class FeaturePartitionData(NamedTuple): 73 | r""" Data and indexing info of a node/edge feature partition. 74 | """ 75 | # node/edge feature tensor 76 | feats: Optional[torch.Tensor] 77 | # node/edge ids tensor corresponding to `feats` 78 | ids: Optional[torch.Tensor] 79 | # feature cache tensor 80 | cache_feats: Optional[torch.Tensor] 81 | # cached node/edge ids tensor corresponding to `cache_feats` 82 | cache_ids: Optional[torch.Tensor] 83 | 84 | HeteroGraphPartitionData = Dict[EdgeType, GraphPartitionData] 85 | HeteroFeaturePartitionData = Dict[Union[NodeType, EdgeType], FeaturePartitionData] 86 | 87 | # Types for neighbor sampling ################################################## 88 | 89 | Seeds = Union[torch.Tensor, str] 90 | InputNodes = Union[Seeds, NodeType, Tuple[NodeType, Seeds], Tuple[NodeType, List[Seeds]]] 91 | EdgeIndexTensor = Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] 92 | InputEdges = Union[EdgeIndexTensor, EdgeType, Tuple[EdgeType, EdgeIndexTensor]] 93 | NumNeighbors = Union[List[int], Dict[EdgeType, List[int]]] 94 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from .common import * 17 | from .device import * 18 | from .exit_status import * 19 | from .mixin import * 20 | from .singleton import * 21 | from .tensor import * 22 | from .topo import * 23 | from .units import * 24 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/device.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import threading 17 | from typing import Optional 18 | 19 | import torch 20 | 21 | 22 | def get_available_device(device: Optional[torch.device] = None) -> torch.device: 23 | r""" Get an available device. If the input device is not ``None``, it will 24 | be returened directly. Otherwise an available device will be choosed ( 25 | current cuda device will be preferred if available). 26 | """ 27 | if device is not None: 28 | return torch.device(device) 29 | if torch.cuda.is_available(): 30 | return torch.device('cuda', torch.cuda.current_device()) 31 | return torch.device('cpu') 32 | 33 | 34 | _cuda_device_assign_lock = threading.RLock() 35 | _cuda_device_rank = 0 36 | 37 | def assign_device(): 38 | r""" Assign an device to use, the cuda device will be preferred if available. 39 | """ 40 | if torch.cuda.is_available(): 41 | global _cuda_device_rank 42 | with _cuda_device_assign_lock: 43 | device_rank = _cuda_device_rank 44 | _cuda_device_rank = (_cuda_device_rank + 1) % torch.cuda.device_count() 45 | return torch.device('cuda', device_rank) 46 | return torch.device('cpu') 47 | 48 | 49 | def ensure_device(device: torch.device): 50 | r""" Make sure that current cuda kernel corresponds to the assigned device. 51 | """ 52 | if (device.type == 'cuda' and 53 | device.index != torch.cuda.current_device()): 54 | torch.cuda.set_device(device) 55 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/exit_status.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import atexit 17 | 18 | 19 | python_exit_status = False 20 | r"""Whether Python is shutting down. This flag is guaranteed to be set before 21 | the Python core library resources are freed, but Python may already be exiting 22 | for some time when this is set. 23 | 24 | Hook to set this flag is `_set_python_exit_flag`, and is same as used in 25 | Pytorch's dataloader: 26 | https://github.com/pytorch/pytorch/blob/f1a6f32b72b7c2b73277f89bbf7e7459a400d80a/torch/utils/data/_utils/__init__.py 27 | """ 28 | 29 | def _set_python_exit_flag(): 30 | global python_exit_status 31 | python_exit_status = True 32 | 33 | atexit.register(_set_python_exit_flag) 34 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/mixin.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | class CastMixin: 17 | r""" This class is same as PyG's :class:`~torch_geometric.utils.CastMixin`: 18 | https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/utils/mixin.py 19 | """ 20 | @classmethod 21 | def cast(cls, *args, **kwargs): 22 | if len(args) == 1 and len(kwargs) == 0: 23 | elem = args[0] 24 | if elem is None: 25 | return None 26 | if isinstance(elem, CastMixin): 27 | return elem 28 | if isinstance(elem, (tuple, list)): 29 | return cls(*elem) 30 | if isinstance(elem, dict): 31 | return cls(**elem) 32 | return cls(*args, **kwargs) 33 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/singleton.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | def singleton(cls): 17 | r""" Singleton class decorator. 18 | """ 19 | instances = {} 20 | def getinstance(): 21 | if cls not in instances: 22 | instances[cls] = cls() 23 | return instances[cls] 24 | return getinstance 25 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/tensor.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from typing import Any, List, Union 17 | 18 | import numpy 19 | import torch 20 | 21 | 22 | def tensor_equal_with_device(lhs: torch.Tensor, rhs: torch.Tensor): 23 | r""" Check whether the data and device of two tensors are same. 24 | """ 25 | if lhs.device == rhs.device: 26 | return torch.equal(lhs, rhs) 27 | return False 28 | 29 | 30 | def id2idx(ids: Union[List[int], torch.Tensor]): 31 | r""" Get tensor of mapping from id to its original index. 32 | """ 33 | if not isinstance(ids, torch.Tensor): 34 | ids = torch.tensor(ids, dtype=torch.int64) 35 | ids = ids.to(torch.int64) 36 | max_id = torch.max(ids).item() 37 | id2idx = torch.zeros(max_id + 1, dtype=torch.int64, device=ids.device) 38 | id2idx[ids] = torch.arange(ids.size(0), dtype=torch.int64, device=ids.device) 39 | return id2idx 40 | 41 | 42 | def convert_to_tensor(data: Any, dtype: torch.dtype = None): 43 | r""" Convert the input data to a tensor based type. 44 | """ 45 | if isinstance(data, dict): 46 | new_data = {} 47 | for k, v in data.items(): 48 | new_data[k] = convert_to_tensor(v, dtype) 49 | return new_data 50 | if isinstance(data, list): 51 | new_data = [] 52 | for v in data: 53 | new_data.append(convert_to_tensor(v, dtype)) 54 | return new_data 55 | if isinstance(data, tuple): 56 | return tuple(convert_to_tensor(list(data), dtype)) 57 | if isinstance(data, torch.Tensor): 58 | return data.type(dtype) if dtype is not None else data 59 | if isinstance(data, numpy.ndarray): 60 | return ( 61 | torch.from_numpy(data).type(dtype) if dtype is not None 62 | else torch.from_numpy(data) 63 | ) 64 | return data 65 | 66 | 67 | def apply_to_all_tensor(data: Any, tensor_method, *args, **kwargs): 68 | r""" Apply the specified method to all tensors contained by the 69 | input data recursively. 70 | """ 71 | if isinstance(data, dict): 72 | new_data = {} 73 | for k, v in data.items(): 74 | new_data[k] = apply_to_all_tensor(v, tensor_method, *args, **kwargs) 75 | return new_data 76 | if isinstance(data, list): 77 | new_data = [] 78 | for v in data: 79 | new_data.append(apply_to_all_tensor(v, tensor_method, *args, **kwargs)) 80 | return new_data 81 | if isinstance(data, tuple): 82 | return tuple(apply_to_all_tensor(list(data), tensor_method, *args, **kwargs)) 83 | if isinstance(data, torch.Tensor): 84 | return tensor_method(data, *args, **kwargs) 85 | return data 86 | 87 | 88 | def share_memory(data: Any): 89 | r""" Share memory for all tensors contained by the input data. 90 | """ 91 | return apply_to_all_tensor(data, torch.Tensor.share_memory_) 92 | 93 | 94 | def squeeze(data: Any): 95 | r""" Squeeze all tensors contained by the input data. 96 | """ 97 | return apply_to_all_tensor(data, torch.Tensor.squeeze) 98 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/topo.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | from typing import Optional, Tuple 17 | 18 | import torch 19 | import torch_sparse 20 | 21 | 22 | def ptr2ind(ptr: torch.Tensor) -> torch.Tensor: 23 | r""" Convert an index pointer tensor to an indice tensor. 24 | """ 25 | ind = torch.arange(ptr.numel() - 1, device=ptr.device) 26 | return ind.repeat_interleave(ptr[1:] - ptr[:-1]) 27 | 28 | 29 | def coo_to_csr( 30 | row: torch.Tensor, 31 | col: torch.Tensor, 32 | edge_id: Optional[torch.Tensor] = None, 33 | edge_weight: Optional[torch.Tensor] = None, 34 | node_sizes: Optional[Tuple[int, int]] = None 35 | ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: 36 | r""" Tranform edge index from COO to CSR. 37 | 38 | Args: 39 | row (torch.Tensor): The row indices. 40 | col (torch.Tensor): The column indices. 41 | edge_id (torch.Tensor, optional): The edge ids corresponding to the input 42 | edge index. 43 | edge_weight (torch.Tensor, optional): The edge weights corresponding to the 44 | input edge index. 45 | node_sizes (Tuple[int, int], optional): The number of nodes in row and col. 46 | """ 47 | if node_sizes is None: 48 | node_sizes = (int(row.max()) + 1, int(col.max()) + 1) 49 | assert len(node_sizes) == 2 50 | assert row.numel() == col.numel() 51 | if edge_id is not None: 52 | assert edge_id.numel() == row.numel() 53 | adj_t = torch_sparse.SparseTensor( 54 | row=row, col=col, value=edge_id, sparse_sizes=node_sizes 55 | ) 56 | edge_ids, edge_weights = adj_t.storage.value(), None 57 | 58 | if edge_weight is not None: 59 | assert edge_weight.numel() == row.numel() 60 | adj_w = torch_sparse.SparseTensor( 61 | row=row, col=col, value=edge_weight, sparse_sizes=node_sizes 62 | ) 63 | edge_weights = adj_w.storage.value() 64 | 65 | return adj_t.storage.rowptr(), adj_t.storage.col(), edge_ids, edge_weights 66 | 67 | 68 | def coo_to_csc( 69 | row: torch.Tensor, 70 | col: torch.Tensor, 71 | edge_id: Optional[torch.Tensor] = None, 72 | edge_weight: Optional[torch.Tensor] = None, 73 | node_sizes: Optional[Tuple[int, int]] = None, 74 | ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: 75 | r""" Tranform edge index from COO to CSC. 76 | 77 | Args: 78 | row (torch.Tensor): The row indices. 79 | col (torch.Tensor): The column indices. 80 | edge_id (torch.Tensor, optional): The edge ids corresponding to the input 81 | edge index. 82 | edge_weight (torch.Tensor, optional): The edge weights corresponding to the 83 | input edge index. 84 | node_sizes (Tuple[int, int], optional): The number of nodes in row and col. 85 | """ 86 | if node_sizes is not None: 87 | node_sizes = (node_sizes[1], node_sizes[0]) 88 | r_colptr, r_row, r_edge_id, r_edge_weight = coo_to_csr( 89 | col, row, edge_id, edge_weight, node_sizes 90 | ) 91 | return r_row, r_colptr, r_edge_id, r_edge_weight 92 | -------------------------------------------------------------------------------- /graphlearn_torch/python/utils/units.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | UNITS = { 17 | # 18 | "KB": 2**10, 19 | "MB": 2**20, 20 | "GB": 2**30, 21 | # 22 | "K": 2**10, 23 | "M": 2**20, 24 | "G": 2**30, 25 | } 26 | 27 | def parse_size(sz) -> int: 28 | if isinstance(sz, int): 29 | return sz 30 | if isinstance(sz, float): 31 | return int(sz) 32 | if isinstance(sz, str): 33 | for suf, u in sorted(UNITS.items()): 34 | if sz.upper().endswith(suf): 35 | return int(float(sz[:-len(suf)]) * u) 36 | raise Exception(f"invalid size: {sz}") 37 | -------------------------------------------------------------------------------- /graphlearn_torch/v6d/vineyard_utils.h: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include 17 | 18 | #include 19 | #include 20 | #include 21 | 22 | namespace graphlearn_torch { 23 | namespace vineyard_utils { 24 | 25 | using GraphType = vineyard::ArrowFragment; 27 | using vertex_t = GraphType::vertex_t; 28 | using vid_t = vineyard::property_graph_types::VID_TYPE; 29 | using fid_t = int64_t; 30 | 31 | std::tuple ToCSR( 32 | const std::string& ipc_socket, const std::string& object_id_str, 33 | const std::string& v_label_name, const std::string& e_label_id_name, 34 | const std::string& edge_dir, bool has_eid); 35 | 36 | torch::Tensor LoadVertexFeatures( 37 | const std::string& ipc_socket, const std::string& object_id_str, 38 | const std::string& v_label_name, std::vector& vcols 39 | ); 40 | 41 | torch::Tensor LoadEdgeFeatures( 42 | const std::string& ipc_socket, const std::string& object_id_str, 43 | const std::string& e_label_name, std::vector& ecols 44 | ); 45 | 46 | int64_t GetFragVertexOffset( 47 | const std::string& ipc_socket, const std::string& object_id_str, 48 | const std::string& v_label_name 49 | ); 50 | 51 | uint32_t GetFragVertexNum( 52 | const std::string& ipc_socket, const std::string& object_id_str, 53 | const std::string& v_label_name 54 | ); 55 | 56 | 57 | class VineyardFragHandle { 58 | public: 59 | explicit VineyardFragHandle( 60 | const std::string& ipc_socket, const std::string& object_id_str); 61 | torch::Tensor GetFidFromGid(const std::vector& gids); 62 | torch::Tensor GetInnerVertices(const std::string& v_label_name); 63 | private: 64 | std::shared_ptr frag_; 65 | }; 66 | 67 | } // namespace vineyard_utils 68 | } // namespace graphlearn_torch 69 | -------------------------------------------------------------------------------- /install_dependencies.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | root_dir=$(pwd) 6 | third_party_dir=${root_dir}/third_party 7 | # googletest 8 | echo "prepare googletest library ..." 9 | if [ ! -d "${third_party_dir}/googletest/build" ]; then 10 | cd "${third_party_dir}/googletest" 11 | bash build.sh 12 | fi 13 | echo "googletest done." -------------------------------------------------------------------------------- /scripts/build_wheel.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | GLT_ROOT_DIR=$(dirname $(dirname "$(realpath "$0")")) 6 | CORES=$(cat < /proc/cpuinfo | grep -c "processor") 7 | 8 | cd $GLT_ROOT_DIR 9 | 10 | set -x 11 | 12 | bash install_dependencies.sh 13 | python3 -m pip install ninja parameterized 14 | 15 | if [ -z "$WITH_CUDA" ]; then 16 | CUDA_OPTION=OFF 17 | else 18 | CUDA_OPTION=$WITH_CUDA 19 | fi 20 | 21 | if [ -z "$WITH_VINEYARD" ]; then 22 | VINEYARD_OPTION=OFF 23 | else 24 | VINEYARD_OPTION=$WITH_VINEYARD 25 | fi 26 | 27 | # TODO(hongyi): build cpp with v6d 28 | # cmake -DWITH_CUDA=$CUDA_OPTION -DWITH_VINEYARD=$VINEYARD_OPTION . 29 | cmake -DWITH_CUDA=$CUDA_OPTION . 30 | make -j$CORES 31 | 32 | python3 setup.py bdist_wheel 33 | python3 -m pip install dist/* --force-reinstall -------------------------------------------------------------------------------- /scripts/launch_container.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e 4 | 5 | GLT_ROOT_DIR=$(dirname $(dirname "$(realpath "$0")")) 6 | IMAGE_NAME=$1 7 | JOB_NAME=$2 8 | DESTDIR=$3 9 | 10 | WITH_VINEYARD=${WITH_VINEYARD:-OFF} 11 | WITH_CUDA=${WITH_CUDA:-ON} 12 | MOUNT_VOLUME=${MOUNT_VOLUME:-TRUE} 13 | 14 | docker_args="-itd --name $JOB_NAME --shm-size=1g -e WITH_VINEYARD=$WITH_VINEYARD -e WITH_CUDA=$WITH_CUDA" 15 | 16 | if [ "$WITH_CUDA" = "ON" ] 17 | then 18 | docker_args+=" --gpus all" 19 | fi 20 | 21 | if [ "$MOUNT_VOLUME" = "TRUE" ] 22 | then 23 | # Option 1: Using -v option 24 | docker_args+=" -v $GLT_ROOT_DIR:$DESTDIR" 25 | docker run $docker_args $IMAGE_NAME bash 26 | else 27 | # Option 2: Using docker cp 28 | docker run $docker_args $IMAGE_NAME bash 29 | docker cp $GLT_ROOT_DIR $JOB_NAME:$DESTDIR 30 | fi -------------------------------------------------------------------------------- /scripts/remove_container.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e 4 | 5 | JOBNAME=$1 6 | 7 | is_exist=0 8 | docker inspect $JOBNAME > /dev/null || is_exist=$? 9 | 10 | if [ $is_exist -eq 0 ] 11 | then 12 | docker stop $JOBNAME 13 | docker rm $JOBNAME 14 | fi 15 | 16 | exit 0 17 | -------------------------------------------------------------------------------- /scripts/run_cpp_ut.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | GLT_ROOT_DIR=$(dirname $(dirname "$(realpath "$0")")) 6 | 7 | echo "Running cpp unit tests ..." 8 | 9 | cd $GLT_ROOT_DIR/built/bin/ 10 | 11 | for i in `ls test_*` 12 | do ./$i 13 | done 14 | -------------------------------------------------------------------------------- /scripts/run_python_ut.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | GLT_ROOT_DIR=$(dirname $(dirname "$(realpath "$0")")) 6 | 7 | PYTHON=python 8 | 9 | echo "Running python unit tests ..." 10 | 11 | for file in $(ls -ld $(find $GLT_ROOT_DIR/test/python)) 12 | do 13 | if [[ $file == */test_*.py ]] 14 | then 15 | echo $file 16 | ${PYTHON} $file 17 | sleep 1s 18 | fi 19 | done 20 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import os 17 | import sys 18 | from torch.utils.cpp_extension import BuildExtension 19 | from setuptools import setup 20 | import torch 21 | import subprocess 22 | import re 23 | 24 | # This version string should be updated when releasing a new version. 25 | _VERSION = '0.2.4' 26 | 27 | RELEASE = os.getenv("RELEASE", "FALSE") 28 | ROOT_PATH = os.path.abspath(os.path.join(os.getcwd())) 29 | WITH_VINEYARD = os.getenv('WITH_VINEYARD', 'OFF') 30 | WITH_CUDA = os.getenv('WITH_CUDA', 'ON') 31 | 32 | sys.path.append(os.path.join(ROOT_PATH, 'graphlearn_torch', 'python', 'utils')) 33 | from build_glt import glt_ext_module, glt_v6d_ext_module 34 | 35 | GLT_V6D_EXT_NAME = "py_graphlearn_torch_vineyard" 36 | GLT_EXT_NAME = "py_graphlearn_torch" 37 | 38 | def get_gcc_use_cxx_abi(): 39 | output = subprocess.run("cmake .", capture_output=True, text=True, shell=True) 40 | print('output', str(output)) 41 | match = re.search(r"GCC_USE_CXX11_ABI: (\d)", str(output)) 42 | if match: 43 | return match.group(1) 44 | else: 45 | return None 46 | 47 | GCC_USE_CXX11_ABI = get_gcc_use_cxx_abi() 48 | 49 | class CustomizedBuildExtension(BuildExtension): 50 | def _add_gnu_cpp_abi_flag(self, extension): 51 | gcc_use_cxx_abi = GCC_USE_CXX11_ABI if extension.name == GLT_V6D_EXT_NAME else str(int(torch._C._GLIBCXX_USE_CXX11_ABI)) 52 | print('GCC_USE_CXX11_ABI for {}: {}', extension.name, gcc_use_cxx_abi) 53 | self._add_compile_flag(extension, '-D_GLIBCXX_USE_CXX11_ABI=' + gcc_use_cxx_abi) 54 | 55 | 56 | ext_modules = [ 57 | glt_ext_module( 58 | name=GLT_EXT_NAME, 59 | root_path=ROOT_PATH, 60 | with_cuda=WITH_CUDA == "ON", 61 | release=RELEASE == "TRUE" 62 | ) 63 | ] 64 | 65 | if WITH_VINEYARD == "ON": 66 | ext_modules.append( 67 | glt_v6d_ext_module( 68 | name=GLT_V6D_EXT_NAME, 69 | root_path=ROOT_PATH, 70 | ), 71 | ) 72 | 73 | setup( 74 | name='graphlearn-torch', 75 | version=_VERSION, 76 | author='GLT Team', 77 | description='Graph Learning for PyTorch (GraphLearn-for-PyTorch)', 78 | url="https://github.com/alibaba/graphlearn-for-pytorch", 79 | python_requires='>=3.6', 80 | requires=['torch'], 81 | cmdclass={'build_ext': CustomizedBuildExtension}, 82 | ext_package='graphlearn_torch', 83 | ext_modules=ext_modules, 84 | package_dir={'graphlearn_torch': 'graphlearn_torch/python'}, 85 | packages=[ 86 | 'graphlearn_torch', 'graphlearn_torch.channel', 'graphlearn_torch.data', 87 | 'graphlearn_torch.distributed', 'graphlearn_torch.loader', 88 | 'graphlearn_torch.partition', 'graphlearn_torch.sampler', 89 | 'graphlearn_torch.utils' 90 | ] 91 | ) 92 | -------------------------------------------------------------------------------- /test/cpp/test_tensor_map_serializer.cu: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include "gtest/gtest.h" 17 | 18 | #include "graphlearn_torch/include/common.cuh" 19 | #include "graphlearn_torch/include/tensor_map.h" 20 | 21 | using namespace graphlearn_torch; 22 | 23 | TEST(TensorMapSerializer, Serialization) { 24 | int current_device; 25 | cudaGetDevice(¤t_device); 26 | 27 | TensorMap orignal_map; 28 | orignal_map.emplace("t1", torch::ones({3, 4}, torch::TensorOptions() 29 | .dtype(torch::kFloat) 30 | .device(torch::kCUDA, current_device))); 31 | CUDACheckError(); 32 | orignal_map.emplace("t2", torch::empty({1}, torch::TensorOptions() 33 | .dtype(torch::kChar) 34 | .device(torch::kCPU))); 35 | orignal_map.emplace("t3", torch::zeros({2, 2, 2}, torch::TensorOptions() 36 | .dtype(torch::kInt) 37 | .device(torch::kCPU))); 38 | 39 | auto buf_size = TensorMapSerializer::GetSerializedSize(orignal_map); 40 | char buf[buf_size]; 41 | TensorMapSerializer::Serialize(orignal_map, buf); 42 | 43 | auto loaded_map = TensorMapSerializer::Load(buf); 44 | 45 | auto& t1 = loaded_map.at("t1"); 46 | EXPECT_EQ(t1.scalar_type(), torch::kFloat); 47 | EXPECT_EQ(t1.device().type(), torch::kCPU); 48 | std::vector t1_size_vec{3, 4}; 49 | torch::IntArrayRef t1_sizes(t1_size_vec); 50 | EXPECT_EQ(t1.sizes(), t1_sizes); 51 | for (size_t i = 0; i < 12; i++) { 52 | EXPECT_EQ(t1.data_ptr()[i], 1.0); 53 | } 54 | 55 | auto& t2 = loaded_map.at("t2"); 56 | EXPECT_EQ(t2.scalar_type(), torch::kChar); 57 | EXPECT_EQ(t2.device().type(), torch::kCPU); 58 | EXPECT_EQ(t2.sizes(), torch::IntArrayRef{1}); 59 | EXPECT_EQ(t2.nbytes(), 1); 60 | 61 | auto& t3 = loaded_map.at("t3"); 62 | EXPECT_EQ(t3.scalar_type(), torch::kInt); 63 | EXPECT_EQ(t3.device().type(), torch::kCPU); 64 | std::vector t3_size_vec{2, 2, 2}; 65 | torch::IntArrayRef t3_sizes(t3_size_vec); 66 | EXPECT_EQ(t3.sizes(), t3_sizes); 67 | for (size_t i = 0; i < 8; i++) { 68 | EXPECT_EQ(t3.data_ptr()[i], 0); 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /test/cpp/test_unified_tensor.cu: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include 17 | #include 18 | #include 19 | 20 | #include 21 | #include 22 | #include 23 | 24 | #include "gtest/gtest.h" 25 | 26 | #include "graphlearn_torch/include/common.cuh" 27 | #include "graphlearn_torch/include/unified_tensor.cuh" 28 | 29 | using namespace graphlearn_torch; 30 | 31 | class UnifiedTensorTest : public ::testing::Test { 32 | protected: 33 | void SetUp() override { 34 | auto torch_cpu_tensor1 = torch::ones({128, 128}, torch::kFloat32); 35 | auto torch_cpu_tensor2 = torch_cpu_tensor1 * 2; 36 | 37 | tensors.resize(2); 38 | tensor_devices.resize(2); 39 | tensors[0] = torch_cpu_tensor1; 40 | tensors[1] = torch_cpu_tensor2; 41 | tensor_devices[0] = 0; 42 | tensor_devices[1] = -1; 43 | } 44 | 45 | protected: 46 | std::vector tensors; 47 | std::vector tensor_devices; 48 | }; 49 | 50 | TEST_F(UnifiedTensorTest, CUDA) { 51 | auto unified_tensor = UnifiedTensor(0); // gpu0 52 | std::cout << "Start Init UnifiedTensor from torch cpu tensors." << std::endl; 53 | unified_tensor.InitFrom(tensors, tensor_devices); 54 | std::cout << "Init UnifiedTensor success." << std::endl; 55 | EXPECT_EQ(unified_tensor.Shape(), std::vector({128*2, 128})); 56 | 57 | auto options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA, 0); 58 | auto indices = torch::tensor({10, 20, 200, 210}, options); 59 | std::cout << "Input Indices:" << indices << std::endl; 60 | auto res = unified_tensor[indices]; 61 | CUDACheckError(); 62 | cudaDeviceSynchronize(); 63 | std::cout << "UnifiedTensor on GPU 0 lookups tensor Success." << std::endl; 64 | EXPECT_EQ(res.device().type(), torch::kCUDA); 65 | EXPECT_EQ(res.device().index(), 0); 66 | EXPECT_EQ(res.size(0), 4); 67 | EXPECT_EQ(res.size(1), 128); 68 | std::cout << "Output Results:" << res << std::endl; 69 | } 70 | -------------------------------------------------------------------------------- /test/cpp/test_vineyard.cu: -------------------------------------------------------------------------------- 1 | /* Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | #include 17 | #include "gtest/gtest.h" 18 | 19 | #ifdef WITH_VINEYARD 20 | 21 | #include "graphlearn_torch/include/common.cuh" 22 | #include "graphlearn_torch/include/graph.h" 23 | #include "graphlearn_torch/v6d/vineyard_utils.h" 24 | 25 | #include 26 | #include 27 | 28 | using namespace graphlearn_torch; 29 | 30 | #endif 31 | 32 | 33 | int main(int argc, const char** argv) { 34 | #ifdef WITH_VINEYARD 35 | 36 | // This test case is for the dataset 37 | // https://github.com/GraphScope/gstest/tree/master/modern_graph 38 | 39 | torch::Tensor indptr; 40 | torch::Tensor indices; 41 | torch::Tensor edge_ids; 42 | 43 | grape::InitMPIComm(); 44 | { 45 | grape::CommSpec comm_spec; 46 | comm_spec.Init(MPI_COMM_WORLD); 47 | std::string ipc_socket = std::string(argv[1]); 48 | std::string object_id_str = std::string(argv[2 + comm_spec.fid()]); 49 | 50 | auto g = Graph(); 51 | std::tie(indptr, indices, edge_ids) = ToCSR(ipc_socket, object_id_str, 0, 1, true); 52 | 53 | std::vector vcols = {"age", "id"}; 54 | std::vector ecols = {"weight"}; 55 | auto vfeat = LoadVertexFeatures( 56 | ipc_socket, object_id_str, 0, vcols, graphlearn_torch::DataType::Int64); 57 | auto efeat = LoadEdgeFeatures( 58 | ipc_socket, object_id_str, 0, ecols, graphlearn_torch::DataType::Float64); 59 | EXPECT_EQ(vfeat.size(0), 4); 60 | EXPECT_EQ(efeat.size(0), 2); 61 | EXPECT_EQ(vfeat[0][0].item(), 27); 62 | EXPECT_EQ(vfeat[3][1].item(), 1); 63 | EXPECT_NEAR(efeat[1].item(), 1.0, 1e-6); 64 | 65 | const auto p = edge_ids.data_ptr(); 66 | 67 | EXPECT_EQ(p[0] + p[1] + p[2] + p[3], 6); 68 | 69 | int64_t row_count = 4; 70 | int64_t col_count = 2; 71 | int64_t edge_count = 4; 72 | 73 | g.InitCUDAGraphFromCSR(indptr, indices, 0, GraphMode::DMA); 74 | CUDACheckError(); 75 | cudaDeviceSynchronize(); 76 | 77 | const auto cu_row_ptr = g.GetRowPtr(); 78 | const auto cu_col_idx = g.GetColIdx(); 79 | 80 | int64_t cpu_row_ptr[row_count+1]; 81 | int64_t cpu_col_idx[edge_count]; 82 | cudaMemcpy((void*)cpu_row_ptr, (void*)cu_row_ptr, 83 | sizeof(int64_t) * (row_count+1), cudaMemcpyDeviceToHost); 84 | cudaMemcpy((void*)cpu_col_idx, (void *)cu_col_idx, 85 | sizeof(int64_t) * edge_count, cudaMemcpyDeviceToHost); 86 | EXPECT_NE(cu_row_ptr, nullptr); 87 | 88 | EXPECT_EQ(g.GetRowCount(), row_count); 89 | EXPECT_EQ(g.GetColCount(), col_count); 90 | EXPECT_EQ(g.GetEdgeCount(), edge_count); 91 | for (int32_t i = 0; i < indptr.size(0); ++i) { 92 | EXPECT_EQ(cpu_row_ptr[i], indptr[i].item()); 93 | } 94 | for (int32_t i = 0; i < indices.size(0); ++i) { 95 | EXPECT_EQ(cpu_col_idx[i], indices[i].item()); 96 | } 97 | } 98 | 99 | grape::FinalizeMPIComm(); 100 | 101 | #endif 102 | } 103 | -------------------------------------------------------------------------------- /test/python/test_dist_feature.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import unittest 17 | 18 | import torch 19 | import graphlearn_torch as glt 20 | 21 | 22 | def run_dist_feature_test(world_size: int, rank: int, feature: glt.data.Feature, 23 | partition_book: glt.data.PartitionBook, master_port: int): 24 | glt.distributed.init_worker_group(world_size, rank, 'dist-feature-test') 25 | glt.distributed.init_rpc(master_addr='localhost', master_port=master_port) 26 | 27 | partition2workers = glt.distributed.rpc_sync_data_partitions(world_size, rank) 28 | rpc_router = glt.distributed.RpcDataPartitionRouter(partition2workers) 29 | 30 | current_device = torch.device('cuda', rank % 2) 31 | 32 | dist_feature = glt.distributed.DistFeature( 33 | world_size, rank, feature, partition_book, 34 | local_only=False, rpc_router=rpc_router, 35 | device=current_device 36 | ) 37 | 38 | input = torch.tensor( 39 | [10, 20, 260, 360, 200, 210, 420, 430], 40 | dtype=torch.int64, 41 | device=current_device 42 | ) 43 | expected_features = torch.cat([ 44 | torch.ones(2, 1024, dtype=torch.float32, device=current_device), 45 | torch.zeros(2, 1024, dtype=torch.float32, device=current_device), 46 | torch.ones(2, 1024, dtype=torch.float32, device=current_device)*2, 47 | torch.zeros(2, 1024, dtype=torch.float32, device=current_device) 48 | ]) 49 | res = dist_feature[input] 50 | 51 | tc = unittest.TestCase() 52 | tc.assertTrue(glt.utils.tensor_equal_with_device(res, expected_features)) 53 | 54 | glt.distributed.shutdown_rpc() 55 | 56 | 57 | class DistFeatureTestCase(unittest.TestCase): 58 | def test_dist_feature_lookup(self): 59 | cpu_tensor0 = torch.cat([ 60 | torch.ones(128, 1024, dtype=torch.float32), 61 | torch.ones(128, 1024, dtype=torch.float32)*2 62 | ]) 63 | cpu_tensor1 = torch.cat([ 64 | torch.zeros(128, 1024, dtype=torch.float32), 65 | torch.zeros(128, 1024, dtype=torch.float32) 66 | ]) 67 | 68 | id2index = torch.arange(128 * 4) 69 | id2index[128*2:] -= 128*2 70 | 71 | partition_book = torch.cat([ 72 | torch.zeros(128*2, dtype=torch.long), 73 | torch.ones(128*2, dtype=torch.long) 74 | ]) 75 | partition_book.share_memory_() 76 | 77 | device_group_list = [ 78 | glt.data.DeviceGroup(0, [0]), 79 | glt.data.DeviceGroup(1, [1]) 80 | ] 81 | 82 | split_ratio = 0.8 83 | 84 | feature0 = glt.data.Feature(cpu_tensor0, id2index, 85 | split_ratio, device_group_list) 86 | feature1 = glt.data.Feature(cpu_tensor1, id2index, 87 | split_ratio, device_group_list) 88 | 89 | mp_context = torch.multiprocessing.get_context('spawn') 90 | port = glt.utils.get_free_port() 91 | w0 = mp_context.Process( 92 | target=run_dist_feature_test, 93 | args=(2, 0, feature0, partition_book, port) 94 | ) 95 | w1 = mp_context.Process( 96 | target=run_dist_feature_test, 97 | args=(2, 1, feature1, partition_book, port) 98 | ) 99 | w0.start() 100 | w1.start() 101 | w0.join() 102 | w1.join() 103 | 104 | 105 | if __name__ == "__main__": 106 | unittest.main() 107 | -------------------------------------------------------------------------------- /test/python/test_sample_prob.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import os 17 | import os.path as osp 18 | import unittest 19 | 20 | import torch 21 | import torch_geometric.transforms as T 22 | from torch_geometric.datasets import OGB_MAG 23 | 24 | from graphlearn_torch.data import Topology, Graph 25 | from graphlearn_torch.sampler import NeighborSampler, NodeSamplerInput 26 | 27 | 28 | class SampleProbTestCase(unittest.TestCase): 29 | @unittest.skip("Download too long") 30 | def setUp(self): 31 | path = osp.join(osp.dirname(osp.realpath(__file__)), '../../../data/') 32 | os.system('mkdir '+path+ ' && wget -P '+path+'mag/raw \ 33 | https://graphlearn.oss-cn-hangzhou.aliyuncs.com/data/github/mag.zip \ 34 | && wget -P '+path+ 'mag/raw \ 35 | https://graphlearn.oss-cn-hangzhou.aliyuncs.com/data/github/mag_metapath2vec_emb.zip' 36 | ) 37 | 38 | transform = T.ToUndirected(merge=True) 39 | dataset = OGB_MAG(path, preprocess='metapath2vec', transform=transform) 40 | data = dataset[0] 41 | 42 | # init graphlearn_torch Dataset. 43 | edge_dict, self.node_dict, csr_dict, self.graph_dict = {}, {}, {}, {} 44 | self.req_nums, self.ids = {}, {} 45 | for etype in data.edge_types: 46 | edge_dict[etype] = data[etype]['edge_index'] 47 | 48 | for ntype in data.node_types: 49 | self.node_dict[ntype] = torch.tensor(list(range(len(data[ntype].x)))) 50 | self.input_type = data.node_types[0] 51 | self.ids = torch.randperm(self.node_dict[self.input_type].size(0))[:5] 52 | 53 | for etype, eidx in edge_dict.items(): 54 | csr_dict[etype] = Topology(edge_index=eidx) 55 | self.graph_dict[etype] = Graph(csr_topo=csr_dict[etype]) 56 | self.req_nums[etype] = [5, 5] 57 | # print(f"{etype}: #row={self.graph_dict[etype].row_count} \ 58 | # #edge={self.graph_dict[etype].edge_count} \ 59 | # #col={self.graph_dict[etype].col_count}") 60 | 61 | @unittest.skip("Download too long") 62 | def test_sample_prob(self): 63 | sampler = NeighborSampler(self.graph_dict, self.req_nums) 64 | print("loading done!") 65 | 66 | inputs = NodeSamplerInput( 67 | node=self.ids, 68 | input_type=self.input_type 69 | ) 70 | probs = sampler.sample_prob(inputs, self.node_dict) 71 | 72 | print(probs) 73 | assert(probs['paper'].size(0) == self.node_dict['paper'].size(0)) 74 | assert(probs['author'].size(0) == self.node_dict['author'].size(0)) 75 | assert(probs['field_of_study'].size(0) == 76 | self.node_dict['field_of_study'].size(0)) 77 | assert(probs['institution'].size(0) == self.node_dict['institution'].size(0)) 78 | 79 | 80 | if __name__ == "__main__": 81 | unittest.main() -------------------------------------------------------------------------------- /test/python/test_shm_channel.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import unittest 17 | import torch 18 | import torch.multiprocessing as mp 19 | 20 | from graphlearn_torch.channel import ShmChannel, QueueTimeoutError 21 | 22 | 23 | def run_sender(_, channel): 24 | for i in range(1, 11): 25 | send_map = {} 26 | send_map['from_cpu'] = torch.ones([i, i], dtype=torch.float32) 27 | send_map['from_cuda'] = torch.arange(i, dtype=torch.int32, 28 | device=torch.device('cuda')) 29 | channel.send(send_map) 30 | print("[sender] message {} sent".format(i)) 31 | 32 | 33 | def run_receiver(_, channel): 34 | channel.pin_memory() 35 | print("[receiver] memory pinned!") 36 | tc = unittest.TestCase() 37 | for i in range(1, 11): 38 | received_map = channel.recv() 39 | print("[receiver] message {} received".format(i)) 40 | tc.assertEqual(len(received_map), 2) 41 | tc.assertTrue('from_cpu' in received_map) 42 | tc.assertEqual(received_map['from_cpu'].device, torch.device('cpu')) 43 | tc.assertTrue(torch.equal(received_map['from_cpu'], 44 | torch.ones([i, i], dtype=torch.float32))) 45 | tc.assertTrue('from_cuda' in received_map) 46 | tc.assertEqual(received_map['from_cuda'].device, torch.device('cpu')) 47 | tc.assertTrue(torch.equal(received_map['from_cuda'], 48 | torch.arange(i, dtype=torch.int32))) 49 | try: 50 | channel.recv(10) 51 | except QueueTimeoutError as e: 52 | print('Expected Error', e) 53 | tc.assertTrue(channel.empty()) 54 | 55 | 56 | class SampleQueueCase(unittest.TestCase): 57 | def test_send_and_receive(self): 58 | channel = ShmChannel(capacity=5, shm_size=1024*1024) 59 | ctx1 = mp.spawn(run_sender, args=(channel,), nprocs=1, join=False) 60 | ctx2 = mp.spawn(run_receiver, args=(channel,), nprocs=1, join=False) 61 | ctx1.join() 62 | ctx2.join() 63 | 64 | 65 | if __name__ == "__main__": 66 | unittest.main() 67 | -------------------------------------------------------------------------------- /test/python/test_unified_tensor.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved. 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 | 16 | import unittest 17 | import torch 18 | 19 | from graphlearn_torch.data import UnifiedTensor 20 | from graphlearn_torch.utils import tensor_equal_with_device 21 | 22 | 23 | class UnifiedTensorTestCase(unittest.TestCase): 24 | def test_unified_tensor_with_int_types(self): 25 | self._test_with_dtype(torch.uint8) 26 | self._test_with_dtype(torch.int8) 27 | self._test_with_dtype(torch.int16) 28 | self._test_with_dtype(torch.int32) 29 | self._test_with_dtype(torch.int64) 30 | 31 | def test_unified_tensor_with_float_types(self): 32 | self._test_with_dtype(torch.float16) 33 | self._test_with_dtype(torch.float32) 34 | self._test_with_dtype(torch.float64) 35 | self._test_with_dtype(torch.bfloat16) 36 | self._test_with_dtype(torch.complex64) 37 | self._test_with_dtype(torch.complex128) 38 | 39 | def _test_with_dtype(self, dtype: torch.dtype): 40 | cpu_tensor1 = torch.ones(128, 128, dtype=dtype) 41 | cpu_tensor2 = cpu_tensor1 * 2 42 | 43 | tensors = [cpu_tensor1, cpu_tensor2] 44 | tensor_devices = [0, -1] 45 | 46 | unified_tensor = UnifiedTensor(0, dtype) 47 | unified_tensor.init_from(tensors, tensor_devices) 48 | self.assertEqual(unified_tensor.shape, [128*2, 128]) 49 | input = torch.tensor([10, 20, 200, 210], dtype=torch.int64, 50 | device= torch.device('cuda:0')) 51 | feature = torch.ones(2, 128, dtype=dtype, device=torch.device('cuda:0')) 52 | expected_res = torch.cat((feature, feature*2), 0) 53 | res = unified_tensor[input] 54 | self.assertEqual(expected_res.dtype, res.dtype) 55 | self.assertTrue(tensor_equal_with_device(expected_res, res)) 56 | 57 | 58 | if __name__ == "__main__": 59 | unittest.main() -------------------------------------------------------------------------------- /test/python/vineyard_data/modern_graph/created.csv: -------------------------------------------------------------------------------- 1 | src_id|dst_id|weight|feat0 2 | 1|3|0.4|0.1 3 | 4|5|1.0|0.2 4 | 4|3|0.4|0.3 5 | 6|3|0.2|0.4 -------------------------------------------------------------------------------- /test/python/vineyard_data/modern_graph/knows.csv: -------------------------------------------------------------------------------- 1 | src_id|dst_id|weight|feat0|feat1 2 | 1|2|0.5|0.1|0.2 3 | 1|4|1.0|0.2|0.4 4 | -------------------------------------------------------------------------------- /test/python/vineyard_data/modern_graph/person.csv: -------------------------------------------------------------------------------- 1 | id|name|age|feat0|feat1|label 2 | 2|vadas|27|0.1|0.2|0 3 | 6|peter|35|0.2|0.4|2 4 | 4|josh|32|0.3|0.6|1 5 | 1|marko|29|0.4|0.8|0 6 | -------------------------------------------------------------------------------- /test/python/vineyard_data/modern_graph/software.csv: -------------------------------------------------------------------------------- 1 | id|name|lang|feat0 2 | 3|lop|java|0.1 3 | 5|ripple|java|0.2 4 | -------------------------------------------------------------------------------- /third_party/googletest/build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -eo pipefail 4 | 5 | script_dir=$(dirname "$(realpath "$0")") 6 | code_src=${script_dir}/googletest 7 | install_prefix=${script_dir}/build 8 | 9 | export CXXFLAGS="-fPIC -D_GLIBCXX_USE_CXX11_ABI=0" 10 | 11 | cd "${code_src}" && mkdir -p tmp && \ 12 | cd tmp && cmake -DCMAKE_INSTALL_PREFIX="${install_prefix}" .. && \ 13 | make -j && mkdir -p "${install_prefix}" && make install && cd .. && \ 14 | rm -rf tmp 15 | 16 | unset CXXFLAGS 17 | --------------------------------------------------------------------------------