Project: GraphScope (App Definition and Development / Database)
File: tutorial_node_classification_pyg_k8s.md

This tutorial presents a server-client example that illustrates how GraphScope trains the GraphSAGE model (implemented in PyG) for a node classification task on a Kubernetes cluster. ```python import graphscope as gs from graphscope.dataset import loadogbnarxiv gs.setoption(loglevel="DEBUG") gs.setoption(showlog=True) params = { "NUMSERVERNODES": 2, "NUMCLIENTNODES": 2, } sess = gs.session( with_dataset=True, k8sservicetype="NodePort", k8svineyardmem="8Gi", k8senginemem="8Gi", vineyardsharedmem="8Gi", k8simagepull_policy="IfNotPresent", k8simagetag="0.26.0a20240115-x86_64", numworkers=params["NUMSERVER_NODES"], ) g = loadogbnarxiv(sess=sess, prefix="/dataset/ogbn_arxiv") ``` ```python gltgraph = gs.graphlearntorch( g, edges=[ ("paper", "citation", "paper"), ], node_features={ "paper": [f"feat_{i}" for i in range(128)], }, node_labels={ "paper": "label", }, edge_dir="out", randomnodesplit={ "num_val": 0.1, "num_test": 0.1, }, numclients=params["NUMCLIENT_NODES"], manifest_path="./client.yaml", clientfolderpath="./", ) print("Exiting...") ``` ```yaml apiVersion: "kubeflow.org/v1" kind: PyTorchJob metadata: name: graphlearn-torch-client namespace: default spec: pytorchReplicaSpecs: Master: replicas: 1 restartPolicy: OnFailure template: spec: containers: name: pytorch image: registry.cn-hongkong.aliyuncs.com/graphscope/graphlearn-torch-client:0.26.0a20240115-x86_64 imagePullPolicy: IfNotPresent command: bash -c |- python3 /workspace/client.py --noderank 0 --masteraddr ${MASTERADDR} --numservernodes ${NUMSERVERNODES} --numclientnodes ${NUMCLIENT_NODES} volumeMounts: mountPath: /dev/shm name: cache-volume mountPath: /workspace name: client-volume volumes: name: cache-volume emptyDir: medium: Memory sizeLimit: "8G" name: client-volume configMap: name: graphlearn-torch-client-config Worker: replicas: ${NUMWORKERREPLICAS} restartPolicy: OnFailure template: spec: containers: name: pytorch image: registry.cn-hongkong.aliyuncs.com/graphscope/graphlearn-torch-client:0.26.0a20240115-x86_64 imagePullPolicy: IfNotPresent command: bash -c |- python3 /workspace/client.py --noderank $((${MYPODNAME: -1}+1)) --masteraddr ${MASTERADDR} --groupmaster ${GROUPMASTER} --numservernodes ${NUMSERVERNODES} --numclientnodes ${NUMCLIENT_NODES} env: name: GROUP_MASTER value: graphlearn-torch-client-master-0 name: MYPODNAME valueFrom: fieldRef: fieldPath: metadata.name volumeMounts: mountPath: /dev/shm name: cache-volume mountPath: /workspace name: client-volume volumes: name: cache-volume emptyDir: medium: Memory sizeLimit: "8G" name: client-volume configMap: name: graphlearn-torch-client-config ``` ```python import argparse import time from typing import List import torch import torch.nn.functional as F from torch.distributed.algorithms.join import Join from torch.nn.parallel import DistributedDataParallel from torch_geometric.nn import GraphSAGE import graphscope as gs import graphscope.learning.graphlearn_torch as glt from graphscope.learning.gltorchgraph import GLTorchGraph from graphscope.learning.graphlearn_torch.typing import Split gs.setoption(loglevel="DEBUG") gs.setoption(showlog=True) ``` ```python @torch.no_grad() def test(model, testloader, datasetname): model.eval() xs = [] y_true = [] for i, batch in enumerate(test_loader): if i == 0: device = batch.x.device batch.x = batch.x.to(torch.float32) # TODO x = model.module(batch.x, batch.edgeindex)[: batch.batchsize] xs.append(x.cpu()) ytrue.append(batch.y[: batch.batchsize].cpu()) del batch xs = [t.to(device) for t in xs] ytrue = [t.to(device) for t in ytrue] y_pred = torch.cat(xs, dim=0).argmax(dim=-1, keepdim=True) ytrue = torch.cat(ytrue, dim=0) testacc = sum((ypred.T == ytrue.T)[0]) / len(ytrue.T) return test_acc.item() ``` ```python def runclientproc( glt_graph, group_master: str, num_servers: int, num_clients: int, client_rank: int, serverranklist: List[int], dataset_name: str, epochs: int, batch_size: int, trainingpgmaster_port: int, ): print("-- Initializing client ...") glt.distributed.init_client( numservers=numservers, numclients=numclients, clientrank=clientrank, masteraddr=gltgraph.master_addr, masterport=gltgraph.serverclientmaster_port, numrpcthreads=4, clientgroupname="k8sgltclient", is_dynamic=True, ) currentctx = glt.distributed.getcontext() torch.distributed.initprocessgroup( backend="gloo", rank=current_ctx.rank, worldsize=currentctx.world_size, initmethod="tcp://{}:{}".format(groupmaster, trainingpgmaster_port), ) device = torch.device("cpu") print("-- Creating training dataloader ...") train_loader = glt.distributed.DistNeighborLoader( data=None, num_neighbors=[5, 3, 2], input_nodes=Split.train, batchsize=batchsize, shuffle=True, collect_features=True, to_device=device, worker_options=glt.distributed.RemoteDistSamplingWorkerOptions( serverrank=serverrank_list, num_workers=1,

worker_concurrency=1, buffer_size="256MB", prefetch_size=1, gltgraph=gltgraph, workload_type="train", ), ) print("-- Creating testing dataloader ...") test_loader = glt.distributed.DistNeighborLoader( data=None, num_neighbors=[5, 3, 2], input_nodes=Split.test, batchsize=batchsize, shuffle=False, collect_features=True, to_device=device, worker_options=glt.distributed.RemoteDistSamplingWorkerOptions( serverrank=serverrank_list, num_workers=1, worker_devices=[torch.device("cpu")], worker_concurrency=1, buffer_size="256MB", prefetch_size=1, gltgraph=gltgraph, workload_type="test", ), ) print("-- Initializing model and optimizer ...") model = GraphSAGE( in_channels=128, hidden_channels=128, num_layers=3, out_channels=47, ).to(device) model = DistributedDataParallel(model, device_ids=None) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) print("-- Start training and testing ...") epochs = 10 dataset_name = "ogbn-arxiv" for epoch in range(0, epochs): model.train() start = time.time() with Join([model]): for batch in train_loader: optimizer.zero_grad() batch.x = batch.x.to(torch.float32) # TODO out = model(batch.x, batch.edgeindex)[: batch.batchsize].log_softmax( dim=-1 ) loss = F.nllloss(out, torch.flatten(batch.y[: batch.batchsize])) loss.backward() optimizer.step() end = time.time() print(f"-- Epoch: {epoch:03d}, Loss: {loss:04f} Epoch Time: {end - start}") torch.distributed.barrier() if epoch == 0 or epoch > (epochs // 2): testacc = test(model, testloader, dataset_name) print(f"-- Test Accuracy: {test_acc:.4f}") torch.distributed.barrier() print("-- Shutdowning ...") glt.distributed.shutdown_client() print("-- Exited ...") ``` ```python if name == "main": parser = argparse.ArgumentParser( description="Arguments for distributed training of supervised SAGE with servers." ) parser.add_argument( "--dataset", type=str, default="ogbn-arxiv", help="The name of ogbn arxiv.", ) parser.add_argument( "--numservernodes", type=int, default=2, help="Number of server nodes for remote sampling.", ) parser.add_argument( "--numclientnodes", type=int, default=1, help="Number of client nodes for training.", ) parser.add_argument( "--node_rank", type=int, default=0, help="The node rank of the current role.", ) parser.add_argument( "--epochs", type=int, default=10, help="The number of training epochs. (client option)", ) parser.add_argument( "--batch_size", type=int, default=256, help="Batch size for the training and testing dataloader.", ) parser.add_argument( "--trainingpgmaster_port", type=int, default=9997, help="The port used for PyTorch's process group initialization across all training processes.", ) parser.add_argument( "--trainloadermaster_port", type=int, default=9998, help="The port used for RPC initialization across all sampling workers of training loader.", ) parser.add_argument( "--testloadermaster_port", type=int, default=9999, help="The port used for RPC initialization across all sampling workers of testing loader.", ) parser.add_argument( "--master_addr", type=str, default="localhost", help="The master address of the graphlearn server.", ) parser.add_argument( "--group_master", type=str, default="localhost", help="The master address of the training process group.", ) args = parser.parse_args() print( f" Distributed training example of supervised SAGE with server-client mode. Client {args.node_rank} " ) print(f"* dataset: {args.dataset}") print(f"* total server nodes: {args.numservernodes}") print(f"* total client nodes: {args.numclientnodes}") print(f"* node rank: {args.node_rank}") numservers = args.numserver_nodes numclients = args.numclient_nodes print(f"* epochs: {args.epochs}") print(f"* batch size: {args.batch_size}") print(f"* training process group master port: {args.trainingpgmaster_port}") print(f"* training loader master port: {args.trainloadermaster_port}") print(f"* testing loader master port: {args.testloadermaster_port}") clientrank = args.noderank print(" Loading graph info ...") glt_graph = GLTorchGraph( [ args.master_addr + ":9001", args.master_addr + ":9002", args.master_addr + ":9003", args.master_addr + ":9004", ] ) print(" Launching client processes ...") runclientproc( glt_graph, args.group_master, num_servers, num_clients, client_rank, [serverrank for serverrank in range(num_servers)], args.dataset, args.epochs, args.batch_size, args.trainingpgmaster_port, ) ``` ```shell python3 k8s_launch.py ```