{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Notebook 6: Running Custom Model with GraphStorm CLIs \n", "\n", "Notebook 1 to 5 provides examples about how to use GraphStorm APIs to implement various GNN components and models. These notebooks can run in the GraphStrom Standalone mode, i.e., on a single CPU or GPU of a single machine. To fully leverage GraphStorm's distributed model training and inference capability, however, we need to convert code implemented on these notebook into Python scripts that can be launched with GraphStorm Command Line Interfaces (CLIs).\n", "\n", "This notebook introduces the method of conversion, and explain the key components of the example Python scripts. For this notebook, we use the custom model developed in the [Notebook 4: Use GraphStorm APIs for Customizing Model Components](https://graphstorm.readthedocs.io/en/latest/api/notebooks/Notebook_4_Customized_Models.html) as an example.\n", "\n", "----\n", "\n", "### Prerequisites\n", "\n", "- GraphStorm. Please find [more details on installation of GraphStorm](https://graphstorm.readthedocs.io/en/latest/install/env-setup.html#setup-graphstorm-with-pip-packages).\n", "- ACM data that has been created according to [Notebook 0: Data Preparation](https://graphstorm.readthedocs.io/en/latest/api/notebooks/Notebook_0_Data_Prepare.html), and is stored in the `./acm_gs_1p/` folder." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Brief Introduction and Run CLIs on a Single Machine\n", "\n", "In order to use GraphStorm CLIs, we need to put the custom model into a Python file, which can be called in the [Task-agnostic CLI for model training and inference](https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/single-machine-training-inference.html#task-agnostic-cli-for-model-training-and-inference) as an argument. We build two files for model training and inference separately.\n", "\n", "We can reuse most of the code about the customized `RGAT` module in Notebook 4, , i.e., `Ara_GatLayer`, `Ara_GatEncoder`, and `RgatNCModel`, in the training and inference files.\n", "\n", "For the training file, we can copy and paste the code of the `4.1 Training pipeline` section in Notebook 4, and enclose them in a `fit()` function. Similarly, for the inference file, we can copy and paste the code of the `4.3 Inference pipeline` section in Notebook 4, and enclose them in a `infer()` function.\n", "\n", "We have provided the two files, named `demo_run_train.py` and `demo_run_infer.py` under the [GraphStorm API documentation folder](https://github.com/awslabs/graphstorm/tree/main/docs/source/api/notebooks). With the two files, we can call GraphStorm's task-agnostic CLI to run our custom model as shown below." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# download the example yaml configuration file\n", "!wget -O acm_nc.yaml https://github.com/awslabs/graphstorm/raw/main/examples/use_your_own_data/acm_nc.yaml\n", "\n", "# CLI for the custom RGAT model training\n", "!python -m graphstorm.run.launch \\\n", " --part-config ./acm_gs_1p/acm.json \\\n", " --num-trainers 4 \\\n", " --num-servers 1 \\\n", " --num-samplers 0 \\\n", " demo_run_train.py --cf acm_nc.yaml \\\n", " --save-model-path models/ \\\n", " --node-feat-name paper:feat author:feat subject:feat \\\n", " --num-epochs 5 \\\n", " --rgat-encoder-type ara\n", "\n", "# CLI for the custom RGAT model inference\n", "!python -m graphstorm.run.launch \\\n", " --part-config ./acm_gs_1p/acm.json \\\n", " --num-trainers 4 \\\n", " --num-servers 1 \\\n", " --num-samplers 0 \\\n", " demo_run_infer.py --cf acm_nc.yaml \\\n", " --restore-model-path models/epoch-4 \\\n", " --save-prediction-path predictions/ \\\n", " --save-embed-path embeddings/ \\\n", " --node-feat-name paper:feat author:feat subject:feat \\\n", " --rgat-encoder-type ara" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "## CLI argument processing explanation\n", "Compared to the code in [Notebook 4](https://graphstorm.readthedocs.io/en/latest/api/notebooks/Notebook_4_Customized_Models.html), the majority of modifications in the two Python files is related to how to collect and parse GraphStorm CLI configurations. Unlike hard-coding some variables, e.g., `nfeats_4_modeling`, or setting fix input values, e.g., `label_field='label',` or `encoder_type='ara'`, we will need to provide these values via CLI configurations.\n", "\n", "As shown in the above commands, there are three types of configurations passed to the GraphStorm task-agnostic command.\n", "\n", "- [Launch CLI arguments](https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/configuration-run.html#launch-cli-arguments), which direclty follow the `graphstom.run.launch`.\n", "- [Model training and inference configurations](https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/configuration-run.html#model-training-and-inference-configurations), which are predefined in GraphStorm. These configurations can be put into a yaml file which will be the value of `--cf` argument following the training or inference Python file name. You can also set them as arguments too, which will overwrite the same configurations set in the yaml file.\n", "- Configurations specified for custom modules, which are not predefined in GraphStorm, but are used only for the custom modules should be defined as input arguments of training or inference Python files.\n", "\n", "Below we show the main entrance function of the `demo_run_train.py` file." ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "tags": [] }, "outputs": [], "source": [ "import argparse\n", "from graphstorm.config import get_argument_parser\n", "\n", "......\n", "\n", "if __name__ == '__main__':\n", " # Leverage GraphStorm's argument parser to accept configuratioin yaml file\n", " arg_parser = get_argument_parser()\n", "\n", " # parse all arguments and split GraphStorm's built-in arguments from the custom ones\n", " gs_args, unknown_args = arg_parser.parse_known_args()\n", " print(f'GS arguments: {gs_args}')\n", "\n", " # create a new argument parser dedicated for custom arguments\n", " cust_parser = argparse.ArgumentParser(description=\"Customized Arguments\")\n", " # add custom arguments\n", " cust_parser.add_argument('--rgat-encoder-type', type=str, default=\"ara\")\n", " cust_args = cust_parser.parse_args(unknown_args)\n", " print(f'Customized arguments: {cust_args}')\n", "\n", " # use both argument sets in our main function\n", " fit(gs_args, cust_args)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "GraphStorm's config module provides a `get_argument_parser` method, which can create a argument parser, e.g., `arg_parser`, dedicated to process GraphStorm launch CLI arguments and model training and inference configurations. Using the `parse_known_args()` method, the argument parser can extract all GraphStorm built-in configurations, and also return custom arguments, which can be processed by another argument parse, e.g., the `cust_parser`. We can then pass these arguments to the corresponding methods. Please refer to [get_argument_parser API document](https://graphstorm.readthedocs.io/en/latest/api/generated/graphstorm.config.get_argument_parser.html#graphstorm.config.get_argument_parser) for more details about this method." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## GraphStorm `GSConfig` object explanation\n", "Once obtained these arguments, we can use them to create a `GSConfig` object and then pass the object to different modules to get related configurations. The `GSConfig` object checks every argument's format and value to ensure compliance with GraphStorm specifications. Below cells show the code of creating the `GSConfig` object and examples of how to use it to pass configurations. For example, we can pass the IP list file, GraphStorm backend, and the local rank configurations to GraphStorm distributed context initialization function, i.e., `gs.initialize()`, to start GraphStorm distributed context.\n", "\n", "For more details of `GSConfig`, please refer to the [GSConfig API documentation page](https://graphstorm.readthedocs.io/en/latest/api/generated/graphstorm.config.GSConfig.html#graphstorm.config.GSConfig) ." ] }, { "cell_type": "code", "execution_count": 20, "metadata": { "tags": [] }, "outputs": [], "source": [ "# in demo_run_train.py file\n", "\n", "from graphstorm.config import GSConfig\n", "\n", "......\n", "\n", "def fit(gs_args, cust_args):\n", " # Utilize GraphStorm's GSConfig class to accept arguments\n", " config = GSConfig(gs_args)\n", "\n", " # Initialize distributed training and inference context\n", " gs.initialize(ip_config=config.ip_config, backend=config.backend, local_rank=config.local_rank)\n", " acm_data = gs.dataloading.GSgnnData(part_config=config.part_config)\n", "\n", " ......\n", "\n", " model = RgatNCModel(g=acm_data.g,\n", " num_heads=config.num_heads,\n", " num_hid_layers=config.num_layers,\n", " node_feat_field=config.node_feat_name,\n", " hid_size=config.hidden_size,\n", " num_classes=config.num_classes,\n", " encoder_type=cust_args.rgat_encoder_type) # here use the custom argument instead of GSConfig\n", "\n", " ......" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# in demo_run_infer.py file\n", "\n", "from graphstorm.config import GSConfig\n", "\n", "......\n", "\n", "def infer(gs_args, cust_args):\n", " # Utilize GraphStorm's GSConfig class to accept arguments\n", " config = GSConfig(gs_args)\n", "\n", " ......\n", "\n", " model = RgatNCModel(g=acm_data.g,\n", " num_heads=config.num_heads,\n", " num_hid_layers=config.num_layers,\n", " node_feat_field=config.node_feat_name,\n", " hid_size=config.hidden_size,\n", " num_classes=config.num_classes,\n", " encoder_type=cust_args.rgat_encoder_type) # here use the custom argument instead of GSConfig\n", "\n", " model.restore_model(config.restore_model_path)\n", "\n", " ......" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Run CLIs on a Distributed Cluster \n", "\n", "It is easy to modify the command in the above cell to run them on a [Distributed clusters](https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/distributed/cluster.html). We need conduct three additional operations:\n", "\n", "1. As demonstrated in [User Your Own Data tutorial](https://graphstorm.readthedocs.io/en/latest/tutorials/own-data.html#run-graph-construction), partition the ACM data in multiple partitions, e.g., 2 partitions by setting the argument `--num-parts 2`, and record its JSON file path, e.g., `./acm_gs_2p/acm.json`.\n", "2. Follow the [tutorial of creating a GraphStorm cluster](https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/distributed/cluster.html#create-a-graphstorm-cluster) to prepare a cluster with 2 machines.\n", "3. Prepare an IP list file, e.g., `ip_list.txt` on the cluster, and record its file path, e.g., `./ip_list.txt`.\n", "\n", "Then we just add two addition CLI launch arguments, and run the CLI below on the clusters within a running docker container." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# CLI for the custom RGAT model training\n", "!python -m graphstorm.run.launch \\\n", " --part-config ./acm_gs_2p/acm.json \\\n", " --num-trainers 4 \\\n", " --num-servers 1 \\\n", " --num-samplers 0 \\\n", " --ip-config ./ip_list.txt \\\n", " --ssh-port 2222 \\\n", " demo_run_train.py --cf acm_nc.yaml \\\n", " --save-model-path models/ \\\n", " --node-feat-name paper:feat author:feat subject:feat \\\n", " --num-epochs 5 \\\n", " --rgat-encoder-type ara\n", "\n", "# CLI for the custom RGAT model inference\n", "!python -m graphstorm.run.launch \\\n", " --part-config ./acm_gs_2p/acm.json \\\n", " --num-trainers 4 \\\n", " --num-servers 1 \\\n", " --num-samplers 0 \\\n", " --ip-config ./ip_list.txt \\\n", " --ssh-port 2222 \\\n", " demo_run_infer.py --cf acm_nc.yaml \\\n", " --restore-model-path models/epoch-4 \\\n", " --save-prediction-path predictions/ \\\n", " --save-embed-path embeddings/ \\\n", " --node-feat-name paper:feat author:feat subject:feat \\\n", " --rgat-encoder-type ara" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Run CLIs on an Amazon SageMaker Cluster\n", "\n", "In order to run the custom models on an Amazon SageMaker cluster, we need to conduct four steps:\n", "\n", "1. Partition the ACM data in multiple partitions, e.g., 2 partition, and upload them to an Amazon S3 location, e.g., `s3:///acm_gs_2p`.\n", "2. Upload the configuration yaml file to an Amazon S3 location, e.g., `s3:///acm_nc.yaml`.\n", "3. Git clone [GraphStorm source code](https://github.com/awslabs/graphstorm), and move the `demo_run_train.py` and `demo_run_infer.py` files from the `graphstorm/docs/source/api/notebooks/` folder to the `graphstorm/python/graphstorm/` folder.\n", "4. Follow the [Setup GraphStorm SageMaker Docker Image](https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/distributed/sagemaker.html#step-1-build-a-sagemaker-compatible-docker-image) tutorial to create a docker image.\n", "\n", "Then use the following SageMaker CLIs to run custom model on an Amazon SageMaker cluster. Please refer to the [GraphStorm Model Training and Inference on on SageMaker](https://graphstorm.readthedocs.io/en/latest/cli/model-training-inference/distributed/sagemaker.html#) for more details." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# SageMaker CLIs should be run under the graphstorm/sagemaker folder\n", "!cd //sagemaker/\n", "\n", "# SageMaker CLI for the customized RGAT model training\n", "!python launch/launch_train.py \\\n", " --image-url \\\n", " --region \\\n", " --entry-point run/train_entry.py \\\n", " --role \\\n", " --instance-count 2 \\\n", " --graph-data-s3 s3:///acm_gs_2p \\\n", " --yaml-s3 s3:///acm_nc.yaml \\\n", " --model-artifact-s3 s3:// \\\n", " --graph-name acm \\\n", " --task-type node_classification \\\n", " --custom-script graphstorm/python/graphstorm/demo_run_train.py \\\n", " --node-feat-name paper:feat author:feat subject:feat \\\n", " --num-epochs 5 \\\n", " --rgat-encoder-type ara\n", "\n", "# SageMaker CLI for the customized RGAT model inference\n", "!python launch/launch_infer.py \\\n", " --image-url \\\n", " --region \\\n", " --entry-point run/infer_entry.py \\\n", " --role \\\n", " --instance-count 2 \\\n", " --graph-data-s3 s3:///acm_gs_2p \\\n", " --yaml-s3 s3:///acm_nc.yaml \\\n", " --model-artifact-s3 s3:// \\\n", " --raw-node-mappings-s3 s3:///acm_gs_2p/raw_id_mappings \\\n", " --output-emb-s3 s3:/// \\\n", " --output-prediction-s3 s3:// \\\n", " --graph-name acm \\\n", " --task-type node_classification \\\n", " --custom-script graphstorm/python/graphstorm/demo_run_infer.py \\\n", " --node-feat-name paper:feat author:feat subject:feat \\\n", " --rgat-encoder-type ara" ] } ], "metadata": { "kernelspec": { "display_name": "gsf", "language": "python", "name": "gsf" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.14" } }, "nbformat": 4, "nbformat_minor": 4 }