[Doc] GraphStorm API example notebook 6 converting API examples to run with GraphStorm CLIs

docs/source/api/notebooks/Notebook_6_API2CLI.ipynb

@@ -0,0 +1,336 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Notebook 6: Converting Customized Model Notebooks to Using 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. GraphStorm Standalone mode is good for quick model development and debugging on a relatively small graph dataset. To fully leverage GraphStorm's distributed model training and inference capacity, however, we need to convert code implemented on these notebook into Python scripts that can be launched with GraphStorm Command Line Interfaces (CLIs) that can handle extreme large graphs.\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 customized 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 customized 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 could 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 in the training and inference files, i.e., copy and paste the code of class `Ara_GatLayer`, `Ara_GatEncoder`, and `RgatNCModel` into the two files. \n",
+    "\n",
+    "For the training file, we can copy and paste the code of training pipeline, and enclose them in a `fit()` function. Similarly, for the inference file, we can copy and paste the code of infernece pipeline, and enclose them in a `infer()` function.\n",
+    "\n",
+    "We have created 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). With the two files, we can call GraphStorm's task-agnostic CLIs to run our customized 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 customized 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 customized 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 and become the value of `--cf` argument follow the training or inference Python file. You can also set them as arguments too, which will overwrite the same configurations set in the yaml file.\n",
+    "- Configurations specified for customized modules, which are not predefined by GraphStorm, but are used only for the customized modules. These configurations should be defined as the arguments for training or inference Python files.\n",
+    "\n",
+    "If following the argument placement convension, it is easy to collect and parse them. Below we show the main entrance method 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 customized 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 customized arguments\n",
+    "    cust_parser = argparse.ArgumentParser(description=\"Customized Arguments\")\n",
+    "    # add customized 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 its `parse_known_args()` method, the argument parser can extract all GraphStorm built-in configurations, and also return other customized 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 can check all arguments' format and values to ensure compliance with GraphStorm specifications. Below cells show the code of creating the `GSConfig` object and examples of how to use it. Please refer to the [GSConfig API doc](https://graphstorm.readthedocs.io/en/latest/api/generated/graphstorm.config.GSConfig.html#graphstorm.config.GSConfig) for more details of this class."
+   ]
+  },
+  {
+   "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",
+    "    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 customized 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 customized 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 CLIs 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. Partition the ACM data in multiple partitions, e.g., 2 partition, 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 the 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 CLIs below on the clusters within running docker environment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# CLI for the customized 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 customized 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 customized models on an Amazon SageMaker cluster, we need to conduct four additional operations:\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://<PATH_TO_DATA>/acm_gs_2p`.\n",
+    "2. Upload the configuration yaml file to an Amazon S3 location, e.g., `s3://<PATH_TO_TRAINING_CONFIG>/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. When building the docker image, please make sure the local graphstorm source code have the two Python files moved.\n",
+    "\n",
+    "Then use the following SageMake CLIs to run customized 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 /<path-to-graphstorm>/sagemaker/\n",
+    "\n",
+    "# SageMaker CLI for the customized RGAT model training\n",
+    "!python launch/launch_train.py \\\n",
+    "        --image-url <AMAZON_ECR_IMAGE_URI> \\\n",
+    "        --region <REGION> \\\n",
+    "        --entry-point run/train_entry.py \\\n",
+    "        --role <ROLE_ARN> \\\n",
+    "        --instance-count 2 \\\n",
+    "        --graph-data-s3 s3://<PATH_TO_DATA>/acm_gs_2p \\\n",
+    "        --yaml-s3 s3://<PATH_TO_TRAINING_CONFIG>/acm_nc.yaml \\\n",
+    "        --model-artifact-s3 s3://<PATH_TO_SAVE_TRAINED_MODEL> \\\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 <AMAZON_ECR_IMAGE_URI> \\\n",
+    "        --region <REGION> \\\n",
+    "        --entry-point run/infer_entry.py \\\n",
+    "        --role <ROLE_ARN> \\\n",
+    "        --instance-count 2 \\\n",
+    "        --graph-data-s3 s3://<PATH_TO_DATA>/acm_gs_2p \\\n",
+    "        --yaml-s3 s3://<PATH_TO_TRAINING_CONFIG>/acm_nc.yaml \\\n",
+    "        --model-artifact-s3 s3://<PATH_TO_SAVE_BEST_TRAINED_MODEL> \\\n",
+    "        --raw-node-mappings-s3 s3://<PATH_TO_DATA>/acm_gs_2p/raw_id_mappings \\\n",
+    "        --output-emb-s3 s3://<PATH_TO_SAVE_GENERATED_NODE_EMBEDDING>/ \\\n",
+    "        --output-prediction-s3 s3://<PATH_TO_SAVE_PREDICTION_RESULTS> \\\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
+}