awslabs · GentleZhu · Sep 24, 2023 · Sep 24, 2023 · Sep 25, 2023 · Sep 25, 2023
diff --git a/.github/workflow_scripts/lint_check.sh b/.github/workflow_scripts/lint_check.sh
@@ -4,6 +4,8 @@ cd ../../
 set -ex
 
 python3 -m pip install --upgrade prospector pip
+pip3 uninstall -y astroid
+yes | pip3 install astroid==2.15.7
 FORCE_CUDA=1 python3 -m pip install -e '.[test]'  --no-build-isolation
 pylint --rcfile=./tests/lint/pylintrc ./python/graphstorm/data/*.py
 pylint --rcfile=./tests/lint/pylintrc ./python/graphstorm/dataloading/

diff --git a/docker/build_docker_wholegraph.sh b/docker/build_docker_wholegraph.sh
@@ -0,0 +1,40 @@
+#!/bin/bash
+
+# process argument 1: graphstorm home folder
+if [ -z "$1" ]; then
+    echo "Please provide the graphstorm home folder that the graphstorm codes are cloned to."
+    echo "For example, ./build_docker_wholegraph.sh /graph-storm/"
+    exit 1
+else
+    GSF_HOME="$1"
+fi
+
+# process argument 2: docker image name, default is graphstorm
+if [ -z "$2" ]; then
+    IMAGE_NAME="graphstorm-wholegraph"
+else
+    IMAGE_NAME="$2"
+fi
+
+# process argument 3: image's tag name, default is local
+if [ -z "$3" ]; then
+    TAG="local"
+else
+    TAG="$3"
+fi
+
+# Copy scripts and tools codes to the docker folder
+mkdir -p $GSF_HOME"/docker/code"
+cp -r $GSF_HOME"/python" $GSF_HOME"/docker/code/python"
+cp -r $GSF_HOME"/inference_scripts" $GSF_HOME"/docker/code/inference_scripts"
+cp -r $GSF_HOME"/tools" $GSF_HOME"/docker/code/tools"
+cp -r $GSF_HOME"/training_scripts" $GSF_HOME"/docker/code/training_scripts"
+
+# Build OSS docker for EC2 instances that an pull ECR docker images
+DOCKER_FULLNAME="${IMAGE_NAME}:${TAG}"
+
+echo "Build a local docker image ${DOCKER_FULLNAME}"
+docker build --no-cache -f $GSF_HOME"/docker/wholegraph/Dockerfile" . -t $DOCKER_FULLNAME
+
+# remove the temporary code folder
+rm -rf $GSF_HOME"/docker/code"
diff --git a/docker/wholegraph/Dockerfile b/docker/wholegraph/Dockerfile
@@ -0,0 +1,60 @@
+FROM nvcr.io/nvidia/dgl:23.07-py3
+
+#################################################
+## Install EFA installer
+ARG EFA_INSTALLER_VERSION=latest
+RUN cd $HOME \
+    && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \
+    && tar -xf $HOME/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \
+    && cd aws-efa-installer \
+    && apt-get update \
+    && apt-get install -y libhwloc-dev \
+    && ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify \
+    && rm -rf /var/lib/apt/lists/*
+
+###################################################
+## Install AWS-OFI-NCCL plugin
+ARG AWS_OFI_NCCL_VERSION=v1.7.1-aws
+RUN git clone https://github.com/aws/aws-ofi-nccl.git /opt/aws-ofi-nccl \
+    && cd /opt/aws-ofi-nccl \
+    && git checkout ${AWS_OFI_NCCL_VERSION} \
+    && ./autogen.sh \
+    && ./configure --prefix=/opt/aws-ofi-nccl/ \
+       --with-libfabric=/opt/amazon/efa/ \
+       --with-cuda=/usr/local/cuda \
+    && make && make install
+
+ENV PATH "/opt/amazon/efa/bin:$PATH"
+
+# Install WholeGraph
+COPY wholegraph/install_wholegraph.sh install_wholegraph.sh
+RUN bash install_wholegraph.sh
+
+# Install GraphStorm
+RUN pip install --no-cache-dir boto3 'h5py>=2.10.0' scipy tqdm 'pyarrow>=3' 'transformers==4.28.1' pandas pylint scikit-learn ogb psutil
+RUN git clone https://github.com/awslabs/graphstorm
+
+# Increase nofile limit
+RUN echo "root                soft    nofile          1048576" >> /etc/security/limits.conf \
+    && echo "root                hard    nofile          1048576" >> /etc/security/limits.conf
+
+# Make EFA NCCL plugin the default plugin
+RUN sed -i '/nccl_rdma_sharp_plugin/d' /etc/ld.so.conf.d/hpcx.conf \
+    && echo "/opt/aws-ofi-nccl/lib" >> /etc/ld.so.conf.d/hpcx.conf \
+    && ldconfig
+
+# Set up SSH
+RUN apt-get update && apt-get install -y openssh-client openssh-server && rm -rf /var/lib/apt/lists/*
+ENV SSH_PORT=2222
+RUN cat /etc/ssh/sshd_config > /tmp/sshd_config && \
+    sed "0,/^#Port 22/s//Port ${SSH_PORT}/" /tmp/sshd_config > /etc/ssh/sshd_config
+ENV HOME=/root
+ENV SSHDIR $HOME/.ssh
+RUN mkdir -p ${SSHDIR}
+RUN ssh-keygen -t rsa -f ${SSHDIR}/id_rsa -N ''
+RUN cp ${SSHDIR}/id_rsa.pub ${SSHDIR}/authorized_keys
+RUN touch /root/.ssh/config;echo -e "Host *\n  StrictHostKeyChecking no\n  UserKnownHostsFile=/dev/null\n  Port ${SSH_PORT}" > /root/.ssh/config
+EXPOSE 2222
+RUN mkdir /run/sshd
+
+CMD ["/usr/sbin/sshd", "-D"]
diff --git a/docker/wholegraph/install_wholegraph.sh b/docker/wholegraph/install_wholegraph.sh
@@ -0,0 +1,25 @@
+#!/bin/bash
+git clone https://github.com/fmtlib/fmt.git  /opt/fmt
+cd /opt/fmt
+git checkout 9.1.0
+mkdir build && cd build
+cmake -DCMAKE_POSITION_INDEPENDENT_CODE=TRUE ..
+make
+make install
+
+git clone https://github.com/gabime/spdlog.git  /opt/spdlog
+cd /opt/spdlog && mkdir build && cd build
+cmake .. && make -j
+cp libspdlog.a  /usr/lib/libspdlog.a
+export PYTHON=/usr/bin/python
+
+cd  /opt/rapids/
+git clone https://github.com/rapidsai/wholegraph.git -b branch-23.08
+cd /opt/rapids/wholegraph/
+pip install scikit-build
+export WHOLEGRAPH_CMAKE_CUDA_ARCHITECTURES="70-real;80-real;90"
+# fix a bug in CMakeList.txt when build pylibwholegraph
+old="import sysconfig; print(sysconfig.get_config_var('BINLIBDEST'))"
+string="import sysconfig; print(\"%s/%s\" % (sysconfig.get_config_var(\"LIBDIR\"), sysconfig.get_config_var(\"INSTSONAME\")))"
+sed -i "s|$old|$string|" /opt/rapids/wholegraph/python/pylibwholegraph/CMakeLists.txt
+bash build.sh libwholegraph pylibwholegraph -v
diff --git a/docs/source/_templates/dataloadertemplate.rst b/docs/source/_templates/dataloadertemplate.rst
@@ -0,0 +1,10 @@
+.. role:: hidden
+    :class: hidden-section
+.. currentmodule:: {{ module }}
+
+
+{{ name | underline}}
+
+.. autoclass:: {{ name }}
+    :show-inheritance:
+    :special-members: __iter__, __next__
diff --git a/docs/source/_templates/datasettemplate.rst b/docs/source/_templates/datasettemplate.rst
@@ -0,0 +1,10 @@
+.. role:: hidden
+    :class: hidden-section
+.. currentmodule:: {{ module }}
+
+
+{{ name | underline}}
+
+.. autoclass:: {{ name }}
+    :show-inheritance:
+    :members: prepare_data, get_node_feats, get_edge_feats, get_labels
diff --git a/docs/source/_templates/classtemplate.rst → docs/source/_templates/evaltemplate.rst b/docs/source/_templates/classtemplate.rst → docs/source/_templates/evaltemplate.rst
diff --git a/docs/source/_templates/inferencetemplate.rst b/docs/source/_templates/inferencetemplate.rst
@@ -0,0 +1,11 @@
+.. role:: hidden
+    :class: hidden-section
+.. currentmodule:: {{ module }}
+
+
+{{ name | underline}}
+
+.. autoclass:: {{ name }}
+    :show-inheritance:
+    :members: setup_device, setup_evaluator, evaluator, device, infer
+
diff --git a/docs/source/_templates/modeltemplate.rst b/docs/source/_templates/modeltemplate.rst
@@ -0,0 +1,10 @@
+.. role:: hidden
+    :class: hidden-section
+.. currentmodule:: {{ module }}
+
+
+{{ name | underline}}
+
+.. autoclass:: {{ name }}
+    :show-inheritance:
+    :members: forward, save_model, restore_model, predict, create_optimizer
diff --git a/docs/source/_templates/trainertemplate.rst b/docs/source/_templates/trainertemplate.rst
@@ -0,0 +1,12 @@
+.. role:: hidden
+    :class: hidden-section
+.. currentmodule:: {{ module }}
+
+
+{{ name | underline}}
+
+.. autoclass:: {{ name }}
+    :show-inheritance:
+    :members: setup_device, setup_evaluator, save_model, remove_saved_model, save_topk_models,
+              get_best_model_path, restore_model, evaluator, optimizer, device, fit, eval
+
diff --git a/docs/source/api/graphstorm.customized.rst b/docs/source/api/graphstorm.customized.rst
diff --git a/docs/source/api/graphstorm.dataloading.rst b/docs/source/api/graphstorm.dataloading.rst
@@ -3,29 +3,48 @@
 graphstorm.dataloading
 ==========================
 
-    GraphStorm dataloading module includes a set of graph datasets and dataloaders for different
+    GraphStorm dataloading module includes a set of graph DataSets and DataLoaders for different
     graph machine learning tasks.
 
+    If users would like to customize DataLoaders, please extend those classes in the
+    :ref:`Base DataLoaders <basedataloaders>` section and customize their abstract methods.
+
 .. currentmodule:: graphstorm.dataloading
 
+.. _basedataloaders:
+
+Base DataLoaders
+-------------------
+
+.. autosummary::
+    :toctree: ../generated/
+    :nosignatures:
+    :template: dataloadertemplate.rst
+
+    GSgnnNodeDataLoaderBase
+    GSgnnEdgeDataLoaderBase
+    GSgnnLinkPredictionDataLoaderBase
+
 DataSets
 ------------
+
 .. autosummary::
     :toctree: ../generated/
     :nosignatures:
-    :template: classtemplate.rst
+    :template: datasettemplate.rst
 
     GSgnnNodeTrainData
     GSgnnNodeInferData
     GSgnnEdgeTrainData
     GSgnnEdgeInferData
 
-Dataloaders
+DataLoaders
 ------------
+
 .. autosummary::
     :toctree: ../generated/
     :nosignatures:
-    :template: classtemplate.rst
+    :template: dataloadertemplate.rst
 
     GSgnnNodeDataLoader
     GSgnnEdgeDataLoader

diff --git a/docs/source/api/graphstorm.eval.rst b/docs/source/api/graphstorm.eval.rst
@@ -0,0 +1,38 @@
+.. _apieval:
+
+graphstorm.eval
+=======================
+
+    GraphStorm provides built-in evaluation methods for different Graph Machine
+    Learning (GML) tasks.
+
+    If users want to implement customized evaluators or evaluation methods, a best practice is to
+    extend base evaluators, i.e., the ``GSgnnInstanceEvaluator`` class for node or edge prediction
+    tasks, and ``GSgnnLPEvaluator`` for link prediction tasks, and then implement the abstract methods.
+
+.. currentmodule:: graphstorm.eval
+
+Base Evaluators
+----------------
+
+.. autosummary::
+    :toctree: ../generated/
+    :nosignatures:
+    :template: evaltemplate.rst
+
+    GSgnnInstanceEvaluator
+    GSgnnLPEvaluator
+
+Evaluators
+-----------
+
+.. autosummary::
+    :toctree: ../generated/
+    :nosignatures:
+    :template: evaltemplate.rst
+
+    GSgnnLPEvaluator
+    GSgnnMrrLPEvaluator
+    GSgnnPerEtypeMrrLPEvaluator
+    GSgnnAccEvaluator
+    GSgnnRegressionEvaluator
diff --git a/docs/source/api/graphstorm.evaluator.rst b/docs/source/api/graphstorm.evaluator.rst
diff --git a/docs/source/api/graphstorm.inferrer.rst → docs/source/api/graphstorm.inference.rst b/docs/source/api/graphstorm.inferrer.rst → docs/source/api/graphstorm.inference.rst
@@ -1,6 +1,6 @@
-.. _apiinferrer:
+.. _apiinference:
 
-graphstorm.inferrer
+graphstorm.inference
 ====================
 
     GraphStorm inferrers assemble the distributed inference pipeline for different tasks.
@@ -13,7 +13,7 @@ graphstorm.inferrer
 .. autosummary::
     :toctree: ../generated/
     :nosignatures:
-    :template: classtemplate.rst
+    :template: inferencetemplate.rst
 
     GSgnnLinkPredictionInferrer
     GSgnnNodePredictionInferrer