import os
import torch
import pickle
import dgl
import time
import numpy as np
from memorys import *
from utils import *
from threading import Thread

class DataLoader:
    def __init__(self, data, train_neg_samples, eval_neg_samples, sets, mode, minibatch_parallelism=1, mailbox=None, node_feats=None, edge_feats=None, edge_classification=False):
        self.edge_classification = edge_classification
        if edge_classification:
            train_neg_samples = 0
            eval_neg_samples = 0
            sets = 0
        self.minibatch_parallelism = minibatch_parallelism
        if mode == 'train' and minibatch_parallelism > 1:
            self.path = 'minibatches/{}_{}_{}_{}_{}/'.format(minibatch_parallelism, data, train_neg_samples, eval_neg_samples, sets)
        else:
            self.path = 'minibatches/{}_{}_{}_{}/'.format(data, train_neg_samples, eval_neg_samples, sets)
        if mode != 'train':
            if not os.path.isfile(self.path + 'val_pos_0.pkl'):
                self.path = 'minibatches/{}_{}_eval/'.format(data, eval_neg_samples)
        self.sets = sets
        self.mode = mode
        self.tot_length = len([fn for fn in os.listdir(self.path) if fn.startswith('{}_pos'.format(mode))]) // minibatch_parallelism
        self.idx = 0
        self.rng = np.random.default_rng()
        self.mailbox = mailbox
        self.node_feats = node_feats
        self.edge_feats = edge_feats
        self.cuda_device_id = torch.cuda.current_device()

        self.rank = None
        self.memory_read_buffer = None
        self.mail_read_buffer = None
        self.read_1idx_buffer = None
        self.read_status = None

    def __iter__(self):
        self.idx = 0
        return self

    def __next__(self):
        # only usable in validation/test
        if self.idx < self.tot_length:
            if self.mode == 'train':
                chosen_set = '_' + str(self.rng.integers(self.sets))
            else:
                chosen_set = ''
            with open('{}{}_pos_{}.pkl'.format(self.path, self.mode, self.idx), 'rb') as f:
                pos_mfg = pickle.load(f)
            neg_mfg = None
            if not self.edge_classification:
                with open('{}{}_neg_{}{}.pkl'.format(self.path, self.mode, self.idx, chosen_set), 'rb') as f:
                    neg_mfg = pickle.load(f)
            self.idx += 1
            return pos_mfg, neg_mfg
        else:
            raise StopIteration

    def get(self, idx, num_neg=1, offset=0):
        idx = idx * self.minibatch_parallelism + offset
        with open('{}{}_pos_{}.pkl'.format(self.path, self.mode, idx), 'rb') as f:
            pos_mfg = pickle.load(f)
        neg_mfgs = list()
        for _ in range(num_neg):
            if self.mode == 'train':
                chosen_set = '_' + str(self.rng.integers(self.sets))
            else:
                chosen_set = ''
            if not self.edge_classification:
                with open('{}{}_neg_{}{}.pkl'.format(self.path, self.mode, idx, chosen_set), 'rb') as f:
                    neg_mfgs.append(pickle.load(f))
        return pos_mfg, neg_mfgs

    def _load_and_slice_mfgs(self, pos_idx, neg_idxs, zero_mails=False):
        # only usable in training set
        t_s = time.time()
        with open('{}{}_pos_{}.pkl'.format(self.path, self.mode, pos_idx), 'rb') as f:
            pos_mfg = pickle.load(f)
        if not self.edge_classification:
            for idx, neg_idx in enumerate(neg_idxs):
                with open('{}{}_neg_{}_{}.pkl'.format(self.path, self.mode, pos_idx, neg_idx), 'rb') as f:
                    neg_mfg = pickle.load(f)
                mfg = combine_mfgs(pos_mfg, neg_mfg)
                prepare_input(mfg, self.node_feats, self.edge_feats)
                self.prefetched_mfgs[idx] = mfg
                if zero_mails:
                    num_nodes = mfg.srcdata['ID'].shape[0]
                    mfg.srcdata['mem'] = torch.zeros((num_nodes, self.mailbox.node_memory.shape[1]), device=mfg.srcdata['ID'].device)
                    mfg.srcdata['mail'] = torch.zeros((num_nodes, self.mailbox.mailbox.shape[1]), device=mfg.srcdata['ID'].device)
                else:
                    self.read_1idx_buffer[self.rank][idx][0] = mfg.srcdata['ID'].shape[0]
                    self.read_1idx_buffer[self.rank][idx][1:1 + mfg.srcdata['ID'].shape[0]].copy_(mfg.srcdata['ID'])
        else:
            mfg = pos_mfg
            mfg.combined = False
            prepare_input(mfg, self.node_feats, self.edge_feats)
            self.prefetched_mfgs[0] = mfg
            if zero_mails:
                num_nodes = mfg.srcdata['ID'].shape[0]
                mfg.srcdata['mem'] = torch.zeros((num_nodes, self.mailbox.node_memory.shape[1]), device=mfg.srcdata['ID'].device)
                mfg.srcdata['mail'] = torch.zeros((num_nodes, self.mailbox.mailbox.shape[1]), device=mfg.srcdata['ID'].device)
            else:
                self.read_1idx_buffer[self.rank][0][0] = mfg.srcdata['ID'].shape[0]
                self.read_1idx_buffer[self.rank][0][1:1 + mfg.srcdata['ID'].shape[0]].copy_(mfg.srcdata['ID'])
        if not zero_mails:
            self.read_status[self.rank] = 1
            # print('here2 read_status:', self.read_status.data_ptr(), self.read_status)
        return

    def init_prefetch(self, idx, num_neg=1, offset=0, prefetch_interval=1, rank=None, memory_read_buffer=None, mail_read_buffer=None, read_1idx_buffer=None, read_status=None):
        # initilize and prefetch the first minibatches with all zero node memory and mails
        # only works in training

        self.rank = rank
        self.memory_read_buffer = memory_read_buffer
        self.mail_read_buffer = mail_read_buffer
        self.read_1idx_buffer = read_1idx_buffer
        self.read_status = read_status
        # print('here1 read_status:', read_status.data_ptr(), 'self.read_status:', read_status.data_ptr())

        self.prefetch_interval = prefetch_interval
        self.prefetch_offset = offset
        self.prefetch_num_neg = num_neg
        self.next_prefetch_idx = idx + prefetch_interval

        self.fetched_mfgs = list()
        self.thread_idx = 0
        self.prefetched_mfgs = [None] * num_neg if not self.edge_classification else [None]
        self.prefetch_threads = None
        idx = 0 if idx < 0 else idx
        pos_idx = idx * self.minibatch_parallelism + offset
        neg_idxs = list()
        for i in range(num_neg):
            if not self.edge_classification:
                neg_idx = self.rng.integers(self.sets)
            else:
                neg_idx = None
            neg_idxs.append(neg_idx)
        self.prefetch_thread = Thread(target=self._load_and_slice_mfgs, args=(pos_idx, neg_idxs, True))
        self.prefetch_thread.start()
        return

    def prefetch_next(self):
        # put current prefetched to fetched and start next prefetch
        t_s = time.time()
        self.prefetch_thread.join()
        if not 'mem' in self.prefetched_mfgs[0].srcdata:
            # print('here3 read_status:', self.read_status.data_ptr())
            while self.read_status[self.rank] == 1:
                pass
            if not self.edge_classification:
                for idx, mfg in enumerate(self.prefetched_mfgs):
                    mfg.srcdata['mem'] = self.memory_read_buffer[self.rank][idx][:mfg.srcdata['ID'].shape[0]]
                    mfg.srcdata['mail'] = self.mail_read_buffer[self.rank][idx][:mfg.srcdata['ID'].shape[0]]
            else:
                self.prefetched_mfgs[0].srcdata['mem'] = self.memory_read_buffer[self.rank][0][:self.prefetched_mfgs[0].srcdata['ID'].shape[0]]
                self.prefetched_mfgs[0].srcdata['mail'] = self.mail_read_buffer[self.rank][0][:self.prefetched_mfgs[0].srcdata['ID'].shape[0]]
        self.thread_idx = 0
        self.fetched_mfgs = self.prefetched_mfgs

        # print('\twait for previous thread time {}ms'.format((time.time() - t_s) * 1000))

        if self.next_prefetch_idx != -1:
            if self.next_prefetch_idx >= self.tot_length:
                self.next_prefetch_idx = self.tot_length - 1

            self.prefetched_mfgs = [None] * self.prefetch_num_neg if not self.edge_classification else [None]
            self.prefetch_threads = None

            pos_idx = self.next_prefetch_idx * self.minibatch_parallelism + self.prefetch_offset
            neg_idxs = list()
            for i in range(self.prefetch_num_neg):
                if not self.edge_classification:
                    neg_idx = self.rng.integers(self.sets)
                else:
                    neg_idx = None
                neg_idxs.append(neg_idx)
            # print('launch prefetch thread to prefetch {}'.format(pos_idx))
            self.prefetch_thread = Thread(target=self._load_and_slice_mfgs, args=(pos_idx, neg_idxs))
            self.prefetch_thread.start()

            if self.next_prefetch_idx != self.tot_length - 1:
                self.next_prefetch_idx += self.prefetch_interval
            else:
                self.next_prefetch_idx = -1

    def get_fetched(self):
        ret = self.fetched_mfgs[self.thread_idx]
        self.thread_idx += 1
        return ret

dataloader.py
amazon-science/disttgl
import argparse
import os

parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, help='dataset name')
parser.add_argument('--train_neg_samples', type=int, default=1)
parser.add_argument('--neg_sets', type=int, default=32, help='how many groups of negative samples')
parser.add_argument('--seed', type=int, default=0, help='random seed to use')
parser.add_argument('--gen_eval', action='store_true', help='whether to generate evaluation minibatches')
parser.add_argument('--eval_cap', type=int, default=0, help='maximum number of minibatches in eval and test set')
parser.add_argument('--minibatch_parallelism', type=int, default=1, help='how many nubmer of GPU per minibatches')
parser.add_argument('--edge_classification', action='store_true', help='whether to train with the edge classification task')
args = parser.parse_args()

if args.data in ['GDELT', 'LINK']:
    args.eval_cap = 5000
    args.neg_sets = 16

if args.data in ['GDELT']:
    args.edge_classification = True

import torch
import dgl
import datetime
import random
import math
import threading
import pickle
import numpy as np
from modules import *
from sampler import *
from utils import *
from get_config import *
from tqdm import tqdm
from pathlib import Path

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed(args.seed)

tot_rank = 1
sample_param, memory_param, gnn_param, train_param = get_config(args.data, tot_rank)
if args.train_neg_samples > 0:
    train_param['train_neg_samples'] = args.train_neg_samples

if args.edge_classification:
    args.neg_sets = 0
    train_param['train_neg_samples'] = 0
    train_param['eval_neg_samples'] = 0
    edge_cls = torch.load('DATA/{}/ec_edge_class.pt'.format(args.data))

g, df = load_graph(args.data)
train_edge_end = df[df['ext_roll'].gt(0)].index[0]
val_edge_end = df[df['ext_roll'].gt(1)].index[0]
num_nodes = g['indptr'].shape[0] - 1

sampler = ParallelSampler(g['indptr'], g['indices'], g['eid'], g['ts'].astype(np.float32),
                            sample_param['num_thread'], 1, sample_param['layer'], sample_param['neighbor'],
                            sample_param['strategy']=='recent', sample_param['prop_time'],
                            sample_param['history'], float(sample_param['duration']))
is_bipartite = True if args.data in ['WIKI', 'REDDIT', 'MOOC', 'LASTFM', 'Taobao', 'sTaobao', 'LINK'] else False
if is_bipartite:
    neg_link_sampler = NegLinkSampler(num_nodes, df=df)
else:
    neg_link_sampler = NegLinkSampler(num_nodes)

if not os.path.isdir('minibatches'):
    os.mkdir('minibatches')
if args.minibatch_parallelism == 1:
    path = 'minibatches/{}_{}_{}_{}/'.format(args.data, train_param['train_neg_samples'], train_param['eval_neg_samples'], args.neg_sets)
else:
    path = 'minibatches/{}_{}_{}_{}_{}/'.format(args.minibatch_parallelism, args.data, train_param['train_neg_samples'], train_param['eval_neg_samples'], args.neg_sets)
if not os.path.isdir(path):
    os.mkdir(path)

stats = {'num_nodes': g['indptr'].shape[0] - 1, 'num_edges': len(df)}
with open('minibatches/{}_stats.pkl'.format(args.data), 'wb') as f:
    pickle.dump(stats, f)

train_df = df[:train_edge_end]
val_df = df[train_edge_end:val_edge_end]
test_df = df[val_edge_end:]

mem_ts = torch.zeros(num_nodes)
mail_ts = torch.zeros(num_nodes)
mail_e = torch.zeros(num_nodes, dtype=torch.long) + len(df)

last_updated_mem = None
i = 0
pos_mfgs = list()
pos_mfgs_i = list()
srcs = list()
dsts = list()
tss = list()
eids = list()
for _, rows in tqdm(train_df.groupby(train_df.index // train_param['batch_size']), total=len(train_df) // train_param['batch_size']):
    # positive mfg
    root_nodes = np.concatenate([rows.src.values, rows.dst.values]).astype(np.int32)
    ts = np.tile(rows.time.values, 2).astype(np.float32)
    sampler.sample(root_nodes, ts)
    ret = sampler.get_ret()
    mfg = to_dgl_blocks(ret, sample_param['history'], 0, cuda=False)[0][0]
    mfg.srcdata['ID'] = mfg.srcdata['ID'].long()
    mfg.srcdata['mem_ts'] = mem_ts[mfg.srcdata['ID']]
    mfg.srcdata['mail_ts'] = mail_ts[mfg.srcdata['ID']]
    mfg.srcdata['mail_e'] = mail_e[mfg.srcdata['ID']]
    if args.edge_classification:
        mfg.edge_cls = edge_cls[rows.index.values].float()
    # mfg.start_eidx = eid[0].item()
    # mfg.end_eidx = eid[-1].item() + 1
    pos_mfgs.append(mfg)
    pos_mfgs_i.append(i)
    
    srcs.append(torch.from_numpy(rows.src.values))
    dsts.append(torch.from_numpy(rows.dst.values))
    tss.append(torch.from_numpy(rows.time.values.astype(np.float32)))
    eids.append(torch.from_numpy(rows.index.values))

    # src = torch.from_numpy(rows.src.values)
    # dst = torch.from_numpy(rows.dst.values)
    # ts = torch.from_numpy(ts)
    # eid = torch.from_numpy(rows.index.values)
    # mailseid = torch.cat([eid, eid])
    # mailsts = ts
    # nid = torch.cat([src.unsqueeze(1), dst.unsqueeze(1)], dim=1).reshape(-1)
    # update_mask = torch.zeros(nid.shape[0], dtype=torch.bool)
    # idx = unique_last_idx(nid)
    # nid = torch.cat([src, dst])
    # idx_map = torch.cat([torch.arange(src.shape[0]).unsqueeze(1), torch.arange(dst.shape[0]).unsqueeze(1) + src.shape[0]], dim=1).reshape(-1)
    # idx = idx_map[idx]
    # update_mask[idx] = 1
    # mfg.node_memory_mask = update_mask

    # negative mfg
    if not args.edge_classification:
        for j in range(args.neg_sets):
            root_nodes = neg_link_sampler.sample(len(rows) * train_param['train_neg_samples'])
            ts = np.tile(rows.time.values, train_param['train_neg_samples']).astype(np.float32)
            sampler.sample(root_nodes, ts)
            ret = sampler.get_ret()
            mfg = to_dgl_blocks(ret, sample_param['history'], 0, cuda=False)[0][0]
            mfg.srcdata['ID'] = mfg.srcdata['ID'].long()
            mfg.srcdata['mem_ts'] = mem_ts[mfg.srcdata['ID']]
            mfg.srcdata['mail_ts'] = mail_ts[mfg.srcdata['ID']]
            mfg.srcdata['mail_e'] = mail_e[mfg.srcdata['ID']]
            with open('{}/train_neg_{}_{}.pkl'.format(path, i, j), 'wb') as f:
                pickle.dump(mfg, f)

    if len(pos_mfgs) == args.minibatch_parallelism:

        idx_map = list()
        length = 0
        for src in srcs:
            idx_map.append(torch.cat([torch.arange(src.shape[0]).unsqueeze(1), torch.arange(src.shape[0]).unsqueeze(1) + src.shape[0]], dim=1).reshape(-1) + length)
            length += idx_map[-1].shape[0]
        idx_map = torch.cat(idx_map)

        src = torch.cat(srcs)
        dst = torch.cat(dsts)
        ts = torch.cat(tss)
        eid = torch.cat(eids)

        nid = torch.cat([src.unsqueeze(1), dst.unsqueeze(1)], dim=1).reshape(-1)
        update_mask = torch.zeros(nid.shape[0], dtype=torch.bool)
        idx_raw = unique_last_idx(nid)
        idx = idx_map[idx_raw]
        update_mask[idx] = 1
        start = 0
        for mfg, mfg_i in zip(pos_mfgs, pos_mfgs_i):
            mfg.node_memory_mask = update_mask[start:start + mfg.num_dst_nodes()]
            start += mfg.num_dst_nodes()
            with open('{}/train_pos_{}.pkl'.format(path, mfg_i), 'wb') as f:
                pickle.dump(mfg, f)

        # update mem_ts, mail_ts, and mail_e
        mailseid = torch.cat([eid, eid])
        mailsts = torch.cat([ts, ts])
        nid = torch.cat([src, dst])

        idx_map = torch.cat([torch.arange(src.shape[0]).unsqueeze(1), torch.arange(dst.shape[0]).unsqueeze(1) + src.shape[0]], dim=1).reshape(-1)
        idx = idx_map[idx_raw]

        nid = nid[idx].long()
        mailseid = mailseid[idx]
        mailsts = mailsts[idx]
        mem_ts[nid] = mail_ts[nid]
        mail_ts[nid] = mailsts
        mail_e[nid] = mailseid

        pos_mfgs = list()
        pos_mfgs_i = list()
        srcs = list()
        dsts = list()
        tss = list()
        eids = list()
    i += 1


if args.gen_eval:
    path = 'minibatches/{}_{}_eval/'.format(args.data, train_param['eval_neg_samples'])
    if not os.path.isdir(path):
        os.mkdir(path)

    # validation
    last_updated_mem = None
    i = 0
    for _, rows in tqdm(val_df.groupby(val_df.index // train_param['batch_size']), total=len(val_df) // train_param['batch_size']):
        # positive mfg
        root_nodes = np.concatenate([rows.src.values, rows.dst.values]).astype(np.int32)
        ts = np.tile(rows.time.values, 2).astype(np.float32)
        sampler.sample(root_nodes, ts)
        ret = sampler.get_ret()
        mfg = to_dgl_blocks(ret, sample_param['history'], 0, cuda=False)[0][0]
        mfg.srcdata['ID'] = mfg.srcdata['ID'].long()
        mfg.srcdata['mem_ts'] = mem_ts[mfg.srcdata['ID']]
        mfg.srcdata['mail_ts'] = mail_ts[mfg.srcdata['ID']]
        mfg.srcdata['mail_e'] = mail_e[mfg.srcdata['ID']]
        if args.edge_classification:
            mfg.edge_cls = edge_cls[rows.index.values].float()

        src = torch.from_numpy(rows.src.values)
        dst = torch.from_numpy(rows.dst.values)
        ts = torch.from_numpy(ts)
        eid = torch.from_numpy(rows.index.values)
        mailseid = torch.cat([eid, eid])
        mailsts = ts
        nid = torch.cat([src.unsqueeze(1), dst.unsqueeze(1)], dim=1).reshape(-1)
        update_mask = torch.zeros(nid.shape[0], dtype=torch.bool)
        idx = unique_last_idx(nid)
        nid = torch.cat([src, dst])
        idx_map = torch.cat([torch.arange(src.shape[0]).unsqueeze(1), torch.arange(dst.shape[0]).unsqueeze(1) + src.shape[0]], dim=1).reshape(-1)
        idx = idx_map[idx]

        update_mask[idx] = 1
        mfg.node_memory_mask = update_mask
        mfg.start_eidx = eid[0].item()
        mfg.end_eidx = eid[-1].item() + 1

        with open('{}/val_pos_{}.pkl'.format(path, i), 'wb') as f:
            pickle.dump(mfg, f)

        # negative mfg
        if not args.edge_classification:
            root_nodes = neg_link_sampler.sample(len(rows) * train_param['eval_neg_samples'])
            ts = np.tile(rows.time.values, train_param['eval_neg_samples']).astype(np.float32)
            sampler.sample(root_nodes, ts)
            ret = sampler.get_ret()
            mfg = to_dgl_blocks(ret, sample_param['history'], 0, cuda=False)[0][0]
            mfg.srcdata['ID'] = mfg.srcdata['ID'].long()
            mfg.srcdata['mem_ts'] = mem_ts[mfg.srcdata['ID']]
            mfg.srcdata['mail_ts'] = mail_ts[mfg.srcdata['ID']]
            mfg.srcdata['mail_e'] = mail_e[mfg.srcdata['ID']]
            with open('{}/val_neg_{}.pkl'.format(path, i), 'wb') as f:
                pickle.dump(mfg, f)
            
        # update mem_ts, mail_ts, and mail_e
        nid = nid[idx].long()
        mailseid = mailseid[idx]
        mailsts = mailsts[idx]
        mem_ts[nid] = mail_ts[nid]
        mail_ts[nid] = mailsts
        mail_e[nid] = mailseid
        
        i += 1

        if args.eval_cap > 0:
            if i > args.eval_cap:
                break

    # test
    last_updated_mem = None
    i = 0
    for _, rows in tqdm(test_df.groupby(test_df.index // train_param['batch_size']), total=len(test_df) // train_param['batch_size']):
        # positive mfg
        root_nodes = np.concatenate([rows.src.values, rows.dst.values]).astype(np.int32)
        ts = np.tile(rows.time.values, 2).astype(np.float32)
        sampler.sample(root_nodes, ts)
        ret = sampler.get_ret()
        mfg = to_dgl_blocks(ret, sample_param['history'], 0, cuda=False)[0][0]
        mfg.srcdata['ID'] = mfg.srcdata['ID'].long()
        mfg.srcdata['mem_ts'] = mem_ts[mfg.srcdata['ID']]
        mfg.srcdata['mail_ts'] = mail_ts[mfg.srcdata['ID']]
        mfg.srcdata['mail_e'] = mail_e[mfg.srcdata['ID']]
        if args.edge_classification:
            mfg.edge_cls = edge_cls[rows.index.values].float()
        
        src = torch.from_numpy(rows.src.values)
        dst = torch.from_numpy(rows.dst.values)
        ts = torch.from_numpy(ts)
        eid = torch.from_numpy(rows.index.values)
        mailseid = torch.cat([eid, eid])
        mailsts = ts
        nid = torch.cat([src.unsqueeze(1), dst.unsqueeze(1)], dim=1).reshape(-1)
        update_mask = torch.zeros(nid.shape[0], dtype=torch.bool)
        idx = unique_last_idx(nid)
        nid = torch.cat([src, dst])
        idx_map = torch.cat([torch.arange(src.shape[0]).unsqueeze(1), torch.arange(dst.shape[0]).unsqueeze(1) + src.shape[0]], dim=1).reshape(-1)
        idx = idx_map[idx]

        update_mask[idx] = 1
        mfg.node_memory_mask = update_mask
        mfg.start_eidx = eid[0].item()
        mfg.end_eidx = eid[-1].item() + 1

        with open('{}/test_pos_{}.pkl'.format(path, i), 'wb') as f:
            pickle.dump(mfg, f)

        # negative mfg
        if not args.edge_classification:
            root_nodes = neg_link_sampler.sample(len(rows) * train_param['eval_neg_samples'])
            ts = np.tile(rows.time.values, train_param['eval_neg_samples']).astype(np.float32)
            sampler.sample(root_nodes, ts)
            ret = sampler.get_ret()
            mfg = to_dgl_blocks(ret, sample_param['history'], 0, cuda=False)[0][0]
            mfg.srcdata['ID'] = mfg.srcdata['ID'].long()
            mfg.srcdata['mem_ts'] = mem_ts[mfg.srcdata['ID']]
            mfg.srcdata['mail_ts'] = mail_ts[mfg.srcdata['ID']]
            mfg.srcdata['mail_e'] = mail_e[mfg.srcdata['ID']]
            with open('{}/test_neg_{}.pkl'.format(path, i), 'wb') as f:
                pickle.dump(mfg, f)

        # update mem_ts, mail_ts, and mail_e
        nid = nid[idx].long()
        mailseid = mailseid[idx]
        mailsts = mailsts[idx]
        mem_ts[nid] = mail_ts[nid]
        mail_ts[nid] = mailsts
        mail_e[nid] = mailseid

        i += 1

        if args.eval_cap > 0:
            if i > args.eval_cap:
                break
gen_minibatch.py
amazon-science/disttgl
import math

def get_config(data, tot_rank, minibatch_parallelism=1):
    sample_param = {
        'layer': 1, 
        'neighbor': [10],
        'strategy': 'recent',
        'prop_time': False,
        'history': 1,
        'duration': 0,
        'num_thread': 32
    }
    memory_param = {
        'type': 'node',
        'dim_time': 100,
        'deliver_to': 'self',
        'mail_combine': 'last',
        'memory_update': 'smart',
        'mailbox_size': 1,
        'combine_node_feature': True,
        'dim_out': 100
    }
    gnn_param = {
        'arch': 'transformer_attention',
        'layer': 1,
        'att_head': 2,
        'dim_time': 100,
        'dim_out': 100
    }
    epoch = 10 if data in ['GDELT', 'LINK'] else 100
    train_param = {
        'epoch': math.ceil(epoch / tot_rank * minibatch_parallelism),
        'batch_size': 600,
        'lr': 0.0001 * tot_rank,
        'dropout': 0.2,
        'att_dropout': 0.2
    }

    train_param['train_neg_samples'] = 1
    train_param['eval_neg_samples'] = 49

    if data in ['GDELT', 'LINK']:
        train_param['batch_size'] = 3200
        # train_param['lr'] = 0.00001

    return sample_param, memory_param, gnn_param, train_param
get_config.py
amazon-science/disttgl
import torch
import dgl
import math
import numpy as np

class TimeEncode(torch.nn.Module):

    def __init__(self, dim):
        super(TimeEncode, self).__init__()
        self.dim = dim
        self.w = torch.nn.Linear(1, dim)
        self.w.weight = torch.nn.Parameter((torch.from_numpy(1 / 10 ** np.linspace(0, 9, dim, dtype=np.float32))).reshape(dim, -1))
        self.w.bias = torch.nn.Parameter(torch.zeros(dim))

    def forward(self, t):
        output = torch.cos(self.w(t.reshape((-1, 1))))
        return output

class EdgePredictor(torch.nn.Module):

    def __init__(self, dim_in):
        super(EdgePredictor, self).__init__()
        self.dim_in = dim_in
        self.src_fc = torch.nn.Linear(dim_in, dim_in)
        self.dst_fc = torch.nn.Linear(dim_in, dim_in)
        self.out_fc = torch.nn.Linear(dim_in, 1)

    def forward(self, pos_rst, neg_rst):
        num_edge = pos_rst.shape[0] // 2
        h_src = self.src_fc(pos_rst[:num_edge])
        h_pos_dst = self.dst_fc(pos_rst[num_edge:])
        h_pos_edge = torch.nn.functional.relu(h_src + h_pos_dst)
        prob_pos = self.out_fc(h_pos_edge)
        prob_neg = None
        if neg_rst is not None:
            h_neg_dst = self.dst_fc(neg_rst)
            neg_samples = h_neg_dst.shape[0] // h_src.shape[0]
            h_neg_edge = torch.nn.functional.relu(h_src.tile(neg_samples, 1) + h_neg_dst)
            prob_neg = self.out_fc(h_neg_edge)
        return prob_pos, prob_neg


class TransfomerAttentionLayer(torch.nn.Module):

    def __init__(self, dim_node_feat, dim_edge_feat, dim_time, num_head, dropout, att_dropout, dim_out, combined=False):
        super(TransfomerAttentionLayer, self).__init__()
        self.num_head = num_head
        self.dim_node_feat = dim_node_feat
        self.dim_edge_feat = dim_edge_feat
        self.dim_time = dim_time
        self.dim_out = dim_out
        self.dropout = torch.nn.Dropout(dropout)
        self.att_dropout = torch.nn.Dropout(att_dropout)
        self.att_act = torch.nn.LeakyReLU(0.2)
        self.combined = combined
        if dim_time > 0:
            self.time_enc = TimeEncode(dim_time)
        if combined:
            if dim_node_feat > 0:
                self.w_q_n = torch.nn.Linear(dim_node_feat, dim_out)
                self.w_k_n = torch.nn.Linear(dim_node_feat, dim_out)
                self.w_v_n = torch.nn.Linear(dim_node_feat, dim_out)
            if dim_edge_feat > 0:
                self.w_k_e = torch.nn.Linear(dim_edge_feat, dim_out)
                self.w_v_e = torch.nn.Linear(dim_edge_feat, dim_out)
            if dim_time > 0:
                self.w_q_t = torch.nn.Linear(dim_time, dim_out)
                self.w_k_t = torch.nn.Linear(dim_time, dim_out)
                self.w_v_t = torch.nn.Linear(dim_time, dim_out)
        else:
            if dim_node_feat + dim_time > 0:
                self.w_q = torch.nn.Linear(dim_node_feat + dim_time, dim_out)
            self.w_k = torch.nn.Linear(dim_node_feat + dim_edge_feat + dim_time, dim_out)
            self.w_v = torch.nn.Linear(dim_node_feat + dim_edge_feat + dim_time, dim_out)
        self.w_out = torch.nn.Linear(dim_node_feat + dim_out, dim_out)
        self.layer_norm = torch.nn.LayerNorm(dim_out)

    def forward(self, b):
        if self.dim_time > 0:
            time_feat = self.time_enc(b.edata['dt'])
            zero_time_feat = self.time_enc(torch.zeros(b.num_dst_nodes(), dtype=torch.float32, device=b.src_idx_cuda.device))
        src_data = b.srcdata['h'][b.src_idx_cuda]
        if b.combined:
            q_data = torch.cat([src_data[:b.num_pos_dst], src_data[b.num_pos_idx:b.num_pos_idx + b.num_neg_dst]], dim=0)
            kv_data = torch.cat([src_data[b.num_pos_dst:b.num_pos_idx], src_data[b.num_pos_idx + b.num_neg_dst:]], dim=0)
        else:
            q_data = src_data[:b.num_dst_nodes()]
            kv_data = src_data[b.num_dst_nodes():]
        if b.num_edges() > 0:
            if self.dim_edge_feat == 0:
                Q = self.w_q(torch.cat([q_data, zero_time_feat], dim=1))[b.edges()[1]]
                K = self.w_k(torch.cat([kv_data, time_feat], dim=1))
                V = self.w_v(torch.cat([kv_data, time_feat], dim=1))
            else:
                Q = self.w_q(torch.cat([q_data, zero_time_feat], dim=1))[b.edges()[1]]
                K = self.w_k(torch.cat([kv_data, b.edata['f'], time_feat], dim=1))
                V = self.w_v(torch.cat([kv_data, b.edata['f'], time_feat], dim=1))
            Q = torch.reshape(Q, (Q.shape[0], self.num_head, -1))
            K = torch.reshape(K, (K.shape[0], self.num_head, -1))
            V = torch.reshape(V, (V.shape[0], self.num_head, -1))
            att = dgl.ops.edge_softmax(b, self.att_act(torch.sum(Q*K, dim=2)))
            att = self.att_dropout(att)
            V = torch.reshape(V*att[:, :, None], (V.shape[0], -1))
            b.edata['v'] = V
            b.update_all(dgl.function.copy_e('v', 'm'), dgl.function.sum('m', 'h'))
        else:
            b.dstdata['h'] = torch.zeros((b.num_dst_nodes(), self.dim_out), device=b.src_idx_cuda.device)
        if b.combined:
            rst_pos = torch.cat([b.dstdata['h'][:b.num_pos_dst], src_data[:b.num_pos_dst]], dim=1)
            rst_neg = torch.cat([b.dstdata['h'][b.num_pos_dst:], src_data[b.num_pos_idx:b.num_pos_idx + b.num_neg_dst]], dim=1)
            rst = torch.cat([rst_pos, rst_neg])
        else:
            rst = torch.cat([b.dstdata['h'], src_data[:b.num_dst_nodes()]], dim=1)
        rst = self.w_out(rst)
        rst = torch.nn.functional.relu(self.dropout(rst))
        return self.layer_norm(rst)

class IdentityNormLayer(torch.nn.Module):

    def __init__(self, dim_out):
        super(IdentityNormLayer, self).__init__()
        self.norm = torch.nn.LayerNorm(dim_out)

    def forward(self, b):
        return self.norm(b.srcdata['h'])

class JODIETimeEmbedding(torch.nn.Module):

    def __init__(self, dim_out):
        super(JODIETimeEmbedding, self).__init__()
        self.dim_out = dim_out

        class NormalLinear(torch.nn.Linear):
        # From Jodie code
            def reset_parameters(self):
                stdv = 1. / math.sqrt(self.weight.size(1))
                self.weight.data.normal_(0, stdv)
                if self.bias is not None:
                    self.bias.data.normal_(0, stdv)

        self.time_emb = NormalLinear(1, dim_out)
    
    def forward(self, h, mem_ts, ts):
        time_diff = (ts - mem_ts) / (ts + 1)
        rst = h * (1 + self.time_emb(time_diff.unsqueeze(1)))
        return rst
            
layers.py
amazon-science/disttgl
import os
import time
from multiprocessing import Process

def start_mailbox_daemon(omp_num_threads, group_id, dims_mailbox, edge_classification=False):
    os.environ['OMP_NUM_THREADS'] = str(omp_num_threads)
    os.environ['MKL_NUM_THREADS'] = str(omp_num_threads)

    import torch
    import dgl
    import random
    import numpy as np
    from dgl.utils.shared_mem import create_shared_mem_array, get_shared_mem_array

    num_nodes, mem_dim, max_read_nodes, max_write_nodes, tot_group_rank, mb_tot_group_rank = dims_mailbox
    minibatch_parallelism = tot_group_rank // mb_tot_group_rank

    node_memory = get_shared_mem_array(str(group_id) + 'node_memory', torch.Size([num_nodes, mem_dim]), dtype=torch.float32)
    mails = get_shared_mem_array(str(group_id) + 'mails', torch.Size([num_nodes, 2 * mem_dim]), dtype=torch.float32)

    memory_read_buffer = get_shared_mem_array('{}_memory_read_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes, mem_dim]), dtype=torch.float32)
    mail_read_buffer = get_shared_mem_array('{}_mail_read_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes, 2 * mem_dim]), dtype=torch.float32)
    read_1idx_buffer = get_shared_mem_array('{}_read_1idx_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes + 1]), dtype=torch.int64)
    read_status = get_shared_mem_array('{}_read_status'.format(group_id), torch.Size([tot_group_rank]), dtype=torch.int)
    memory_write_buffer = get_shared_mem_array('{}_memory_write_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes, mem_dim]), dtype=torch.float32)
    mail_write_buffer = get_shared_mem_array('{}_mail_write_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes, 2 * mem_dim]), dtype=torch.float32)
    write_1idx_buffer = get_shared_mem_array('{}_write_1idx_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes + 1]), dtype=torch.int64)
    write_status = get_shared_mem_array('{}_write_status'.format(group_id), torch.Size([tot_group_rank]), dtype=torch.int)
    reset_status = get_shared_mem_array('{}_reset_status'.format(group_id), torch.Size([1]), dtype=torch.int)

    # print('deamon read_status:', read_status.data_ptr(), 'write_status:', write_status.data_ptr())

    def write_rank(rank):
        nonlocal node_memory, mails, memory_write_buffer, mail_write_buffer, write_1idx_buffer
        # print('write for rank {}'.format(rank))
        with torch.no_grad():
            write_1idx = write_1idx_buffer[rank]
            write_idx = write_1idx[1:1 + write_1idx[0]]
            node_memory[write_idx] = memory_write_buffer[rank][:write_idx.shape[0]]
            mails[write_idx] = mail_write_buffer[rank][:write_idx.shape[0]]

    def read_rank(rank):
        nonlocal node_memory, mails, memory_read_buffer, mail_read_buffer, read_1idx_buffer
        # print('read for rank {}'.format(rank))
        with torch.no_grad():
            for i in range(mb_tot_group_rank):
                read_1idx = read_1idx_buffer[rank][i]
                read_idx = read_1idx[1:1 + read_1idx[0]]
                torch.index_select(node_memory, 0, read_idx, out=memory_read_buffer[rank][i][:read_idx.shape[0]])
                torch.index_select(mails, 0, read_idx, out=mail_read_buffer[rank][i][:read_idx.shape[0]])
                if edge_classification:
                    break

    with torch.no_grad():
        while True:
            node_memory.zero_()
            mails.zero_()
            # read_status.zero_()
            # write_status.zero_()
            curr_rank = 0
            reset_flag = False
            while True:
                # write to node memory
                # print('curr_rank:{}'.format(curr_rank))
                for rank in range(curr_rank, curr_rank + minibatch_parallelism):
                    # print('\twait for write request from {}...'.format(rank), end='')
                    while write_status[rank] == 0:
                        if reset_status[0] == 1:
                            reset_flag = True
                            reset_status[0] = 2
                            break
                        # print('\tread status:{} write status:{}'.format(read_status, write_status))
                        # time.sleep(0.01)
                        pass
                    # print('done!')
                    if not reset_flag:
                        write_rank(rank)
                        # print('\tset write_status[{}] to 0'.format(rank))
                        write_status[rank] = 0
                
                if not reset_flag:
                    # serving next rank
                    curr_rank += minibatch_parallelism
                    curr_rank = 0 if curr_rank >= tot_group_rank else curr_rank
                    # print('\t advance curr_rank to {}'.format(curr_rank))
                    # print('\twait for read request from {}...'.format(curr_rank), end='')
                    while not (read_status[curr_rank] == 1):
                        if reset_status[0] == 1:
                            reset_flag = True
                            reset_status[0] = 2
                            break
                        pass
                    if not reset_flag:
                        for rank in range(curr_rank, curr_rank + minibatch_parallelism):
                            while not (read_status[rank] == 1):
                                if reset_status[0] == 1:
                                    reset_flag = True
                                    reset_status[0] = 2
                                    break
                                pass
                            if not reset_flag:
                                # serve read request
                                # print('done!')
                                read_rank(rank)
                                # print('reset read status to 0')
                                read_status[rank] = 0
                if reset_flag:
                    # print('\tbreak inner loop!')
                    break

            
mailbox_daemon.py
amazon-science/disttgl
import torch
import dgl
import time
from layers import TimeEncode

class MailBox():

    def __init__(self, memory_param, num_nodes, dim_edge_feat, _node_memory=None, _mailbox=None):
        self.memory_param = memory_param
        self.dim_edge_feat = dim_edge_feat
        if memory_param['type'] != 'node':
            raise NotImplementedError
        self.node_memory = torch.zeros((num_nodes, memory_param['dim_out']), dtype=torch.float32) if _node_memory is None else _node_memory
        self.mailbox = torch.zeros((num_nodes, 2 * memory_param['dim_out']), dtype=torch.float32) if _mailbox is None else _mailbox
        self.device = torch.device('cpu')
        
    def reset(self):
        self.node_memory.fill_(0)
        self.mailbox.fill_(0)

    def prep_input_mails(self, mfg):
        if hasattr(mfg, 'stream'):
            with torch.cuda.stream(mfg.stream):
                idx = mfg.srcdata['ID'].cpu()
                mfg.srcdata['mem'] = torch.index_select(self.node_memory, 0, idx).to(mfg.srcdata['ID'].device)
                mfg.srcdata['mail'] = torch.index_select(self.mailbox, 0, idx).to(mfg.srcdata['ID'].device)
        else:
            idx = mfg.srcdata['ID'].cpu()
            mfg.srcdata['mem'] = torch.index_select(self.node_memory, 0, idx).to(mfg.srcdata['ID'].device)
            mfg.srcdata['mail'] = torch.index_select(self.mailbox, 0, idx).to(mfg.srcdata['ID'].device)

    def prep_zero_mails(self, mfg):
        num_nodes = mfg.srcdata['ID'].shape[0]
        mfg.srcdata['mem'] = torch.zeros((num_nodes, self.node_memory.shape[1]), device=mfg.srcdata['ID'].device)
        mfg.srcdata['mail'] = torch.zeros((num_nodes, self.mailbox.shape[1]), device=mfg.srcdata['ID'].device)

    def update_memory_and_mailbox(self, mfg, memory):
        with torch.no_grad():
            if mfg.combined:
                mask = mfg.node_memory_mask.to(self.device)
                nid = mfg.srcdata['ID'][:mfg.num_pos_src].to(self.device)[mfg.src_idx[:mask.shape[0]]]
                memory = memory.to(self.device)[mfg.src_idx[:mask.shape[0]]]
            else:
                mask = mfg.node_memory_mask.to(self.device)
                nid = mfg.srcdata['ID'].to(self.device)[mfg.src_idx[:mask.shape[0]]]
                memory = memory.to(self.device)[mfg.src_idx[:mask.shape[0]]]
            masked_nid = nid[mask]
            masked_memory = memory[mask]
            self.node_memory[masked_nid] = masked_memory

            num_root = mask.shape[0] // 2
            mail_src = masked_memory
            mail_dst = torch.cat([memory[num_root:][mask[:num_root]], memory[:num_root][mask[num_root:]]])
            mail = torch.cat([mail_src, mail_dst], dim=1)
            self.mailbox[masked_nid] = mail

class SmartMemoryUpdater(torch.nn.Module):

    def __init__(self, memory_param, dim_in, dim_hid, dim_time, dim_node_feat, num_node, no_learn_node=False):
        super(SmartMemoryUpdater, self).__init__()
        self.dim_hid = dim_hid
        self.dim_node_feat = dim_node_feat
        self.memory_param = memory_param
        self.dim_time = dim_time

        self.updater = torch.nn.GRUCell(dim_in + dim_time, dim_hid)
        self.last_updated_memory = None
        self.last_updated_ts = None
        self.last_updated_nid = None
        if dim_time > 0:
            self.time_enc = TimeEncode(dim_time)
        self.addc = False

        # self.no_learn_node = no_learn_node
        if dim_node_feat > 0:
            if dim_node_feat != dim_hid:
                self.node_feat_map = torch.nn.Linear(dim_node_feat, dim_hid)
        # if not no_learn_node:
        #     if dim_node_feat == 0:
        #         self.constant = torch.nn.Embedding(num_node, dim_hid, sparse=True)

        #     self.choice = torch.nn.Embedding(num_node, 1, sparse=True)
        #     torch.nn.init.constant_(self.choice.weight, 0)

    def forward(self, b, write_buffer=None):
        # torch.cuda.synchronize()
        # t_s=time.time()
        if self.dim_time > 0:
            time_feat = self.time_enc(b.srcdata['mail_ts'].squeeze() - b.srcdata['mem_ts'])
            if 'mail_ef' in b.srcdata:
                b.srcdata['t_mem_input'] = torch.cat([b.srcdata['mail'], b.srcdata['mail_ef'], time_feat], dim=1)
            else:
                b.srcdata['t_mem_input'] = torch.cat([b.srcdata['mail'], time_feat], dim=1)
        updated_memory = self.updater(b.srcdata['t_mem_input'], b.srcdata['mem'])
        new_memory = updated_memory
        # torch.cuda.synchronize()
        # print('t_gru:{:.4f}'.format(time.time() - t_s))

        with torch.no_grad():
            if write_buffer is not None:
                rank = write_buffer.rank
                memory_write_buffer = write_buffer.memory_write_buffer
                mail_write_buffer = write_buffer.mail_write_buffer
                write_1idx_buffer = write_buffer.write_1idx_buffer
                write_status = write_buffer.write_status
                if b.combined:
                    memory = updated_memory[:b.num_pos_src]
                else:
                    memory = updated_memory
                if b.combined:
                    mask = b.node_memory_mask_cuda
                    nid = b.srcdata['ID'][:b.num_pos_src][b.src_idx[:mask.shape[0]]]
                    memory = memory[b.src_idx[:mask.shape[0]]]
                else:
                    mask = b.node_memory_mask_cuda
                    nid = b.srcdata['ID'][b.src_idx[:mask.shape[0]]]
                    memory = memory[b.src_idx[:mask.shape[0]]]

                masked_nid = nid[mask]
                masked_memory = memory[mask]
                write_1idx_buffer[rank][0] = masked_nid.shape[0]
                write_1idx_buffer[rank][1:1 + masked_nid.shape[0]].copy_(masked_nid)
                memory_write_buffer[rank][:masked_nid.shape[0]].copy_(masked_memory)

                num_root = mask.shape[0] // 2
                mail_src = masked_memory
                mail_dst = torch.cat([memory[num_root:][mask[:num_root]], memory[:num_root][mask[num_root:]]])
                mail = torch.cat([mail_src, mail_dst], dim=1)
                mail_write_buffer[rank][:masked_nid.shape[0]].copy_(mail)

                write_status[rank] = 1
                # import pdb; pdb.set_trace()
                # pass
            else:
                if b.combined:
                    self.last_updated_memory = updated_memory[:b.num_pos_src].detach().clone()
                else:
                    self.last_updated_memory = updated_memory.detach().clone()

        # if not self.no_learn_node:
        #     if self.dim_node_feat > 0:
        #         if self.dim_node_feat == self.dim_hid:
        #             constant_memory = b.srcdata['rh']
        #         else:
        #             constant_memory = self.node_feat_map(b.srcdata['rh'])
        #     else:
        #         constant_memory = self.constant(b.srcdata['ID'])
        #     thres = torch.sigmoid(self.choice(b.srcdata['ID']))
        #     b.srcdata['h'] = new_memory * thres + constant_memory * (1 - thres)
        # else:
        if self.dim_node_feat > 0:
            if self.dim_node_feat == self.dim_hid:
                constant_memory = b.srcdata['rh']
            else:
                constant_memory = self.node_feat_map(b.srcdata['rh'])
            b.srcdata['h'] = new_memory + constant_memory
        else:
            b.srcdata['h'] = new_memory
memorys.py
amazon-science/disttgl
import torch
import dgl
import time
from memorys import *
from layers import *

class GeneralModel(torch.nn.Module):

    def __init__(self, dim_node, dim_edge, sample_param, memory_param, gnn_param, train_param, num_node, no_learn_node=False, combined=False, edge_classification=False, edge_classes=0):
        super(GeneralModel, self).__init__()
        self.edge_classification = edge_classification
        self.dim_node = dim_node
        self.dim_node_input = dim_node
        self.dim_edge = dim_edge
        self.sample_param = sample_param
        self.memory_param = memory_param
        if not 'dim_out' in gnn_param:
            gnn_param['dim_out'] = memory_param['dim_out']
        self.gnn_param = gnn_param
        self.train_param = train_param
        if memory_param['type'] == 'node':
            if memory_param['memory_update'] == 'smart':
                self.memory_updater = SmartMemoryUpdater(memory_param, 2 * memory_param['dim_out'] + dim_edge, memory_param['dim_out'], memory_param['dim_time'], dim_node, num_node, no_learn_node=no_learn_node)
            else:
                raise NotImplementedError
            self.dim_node_input = memory_param['dim_out']
        self.layers = torch.nn.ModuleDict()
        if gnn_param['arch'] == 'transformer_attention':
            for h in range(sample_param['history']):
                self.layers['l0h' + str(h)] = TransfomerAttentionLayer(self.dim_node_input, dim_edge, gnn_param['dim_time'], gnn_param['att_head'], train_param['dropout'], train_param['att_dropout'], gnn_param['dim_out'], combined=combined)
            for l in range(1, gnn_param['layer']):
                for h in range(sample_param['history']):
                    self.layers['l' + str(l) + 'h' + str(h)] = TransfomerAttentionLayer(gnn_param['dim_out'], dim_edge, gnn_param['dim_time'], gnn_param['att_head'], train_param['dropout'], train_param['att_dropout'], gnn_param['dim_out'], combined=False)
        else:
            raise NotImplementedError
        if not self.edge_classification:
            self.edge_predictor = EdgePredictor(gnn_param['dim_out'])
        else:
            self.edge_classifier = torch.nn.Linear(2 * gnn_param['dim_out'], edge_classes)
    
    def forward(self, mfg, pos_mfg=None, neg_mfg=None, write_buffer=None):
        # import pdb; pdb.set_trace()
        # torch.cuda.synchronize()
        # t_s=time.time()
        if self.memory_param['type'] == 'node':
            self.memory_updater(mfg, write_buffer=write_buffer)
        # torch.cuda.synchronize()
        # t_mem = time.time() - t_s
        # t_s = time.time()
        rst = self.layers['l0h0'](mfg)
        if mfg.combined:
            pos_rst = rst[:mfg.num_pos_dst]
            neg_rst = rst[mfg.num_pos_dst:]
        else:
            pos_rst = rst
            neg_rst = None
        # if pos_mfg is not None:
        #     pos_mfg.combined = False
        #     neg_mfg.combined = False
        #     self.memory_updater(pos_mfg)
        #     self.memory_updater(neg_mfg)
        #     ref_pos_rst = self.layers['l0h0'](pos_mfg)
        #     ref_neg_rst = self.layers['l0h0'](neg_mfg)
        #     import pdb; pdb.set_trace()
        # torch.cuda.synchronize()
        # t_gat = time.time() - t_s
        # print('t_mem:{:.4f} t_gat:{:.4f}'.format(t_mem, t_gat))
        if not self.edge_classification:
            return self.edge_predictor(pos_rst, neg_rst)
        else:
            rst = torch.cat([pos_rst[:pos_rst.shape[0] // 2], pos_rst[pos_rst.shape[0] // 2:]], dim=1)
            return self.edge_classifier(rst), None

    # def get_emb(self, mfgs):
    #     if self.memory_param['type'] == 'node':
    #         self.memory_updater(mfgs[0])
    #     out = list()
    #     for l in range(self.gnn_param['layer']):
    #         for h in range(self.sample_param['history']):
    #             rst = self.layers['l' + str(l) + 'h' + str(h)](mfgs[l][h])
    #             if l != self.gnn_param['layer'] - 1:
    #                 mfgs[l + 1][h].srcdata['h'] = rst
    #             else:
    #                 out.append(rst)
    #     if self.sample_param['history'] == 1:
    #         out = out[0]
    #     else:
    #         out = torch.stack(out, dim=0)
    #         out = self.combiner(out)[0][-1, :, :]
    #     return out

class NodeClassificationModel(torch.nn.Module):

    def __init__(self, dim_in, dim_hid, num_class):
        super(NodeClassificationModel, self).__init__()
        self.fc1 = torch.nn.Linear(dim_in, dim_hid)
        self.fc2 = torch.nn.Linear(dim_hid, num_class)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.nn.functional.relu(x)
        x = self.fc2(x)
        return x
modules.py
amazon-science/disttgl
import argparse
import yaml
import torch
import time
import random
import numpy as np
import pandas as pd
from tqdm import tqdm
from sampler_core import ParallelSampler, TemporalGraphBlock

class NegLinkSampler:

    def __init__(self, num_nodes, df=None, transduct=False, induct=False, num_rngs=0):
        self.random_generator = None
        self.num_rngs = num_rngs
        self.rng_list = list()
        self.num_nodes = num_nodes
        self.random_sample = True
        self.dst_nodes = None
        if df is not None:
            self.random_sample = False
            train_edge_end = df[df['ext_roll'].gt(0)].index[0]
            if transduct:
                self.dst_nodes = np.unique(df[:train_edge_end].dst.values.astype(np.int64))
            elif induct:
                self.dst_nodes = np.unique(df[train_edge_end:].dst.values.astype(np.int64))
                trans_nodes = np.unique(df[:train_edge_end].dst.values.astype(np.int64))
                self.dst_nodes = self.dst_nodes[np.logical_not(np.isin(self.dst_nodes, trans_nodes))]
            else:
                self.dst_nodes = np.unique(df.dst.values.astype(np.int64))
        self.reset_rng()

    def reset_rng(self, i=None):
        if self.num_rngs > 0:
            if i is None:
                self.random_generator = np.random.default_rng(random.randint(0, self.num_rngs - 1))
            else:
                self.random_generator = np.random.default_rng(i % self.num_rngs)
        else:
            self.random_generator = np.random.default_rng()

    def sample(self, n):
        if self.random_sample:
            return self.random_generator.integers(self.num_nodes, size=n)
        else:
            return self.random_generator.choice(self.dst_nodes, size=n)

if __name__ == '__main__':
    parser=argparse.ArgumentParser()
    parser.add_argument('--data', type=str, help='dataset name')
    parser.add_argument('--config', type=str, help='path to config file')
    parser.add_argument('--batch_size', type=int, default=600, help='path to config file')
    parser.add_argument('--num_thread', type=int, default=64, help='number of thread')
    args=parser.parse_args()

    df = pd.read_csv('DATA/{}/edges.csv'.format(args.data))
    g = np.load('DATA/{}/ext_full.npz'.format(args.data))
    sample_config = yaml.safe_load(open(args.config, 'r'))['sampling'][0]

    sampler = ParallelSampler(g['indptr'], g['indices'], g['eid'], g['ts'].astype(np.float32),
                              args.num_thread, 1, sample_config['layer'], sample_config['neighbor'],
                              sample_config['strategy']=='recent', sample_config['prop_time'],
                              sample_config['history'], float(sample_config['duration']))

    num_nodes = max(int(df['src'].max()), int(df['dst'].max()))
    neg_link_sampler = NegLinkSampler(num_nodes)

    tot_time = 0
    ptr_time = 0
    coo_time = 0
    sea_time = 0
    sam_time = 0
    uni_time = 0
    total_nodes = 0
    unique_nodes = 0
    for _, rows in tqdm(df.groupby(df.index // args.batch_size), total=len(df) // args.batch_size):
        root_nodes = np.concatenate([rows.src.values, rows.dst.values, neg_link_sampler.sample(len(rows))]).astype(np.int32)
        ts = np.concatenate([rows.time.values, rows.time.values, rows.time.values]).astype(np.float32)
        sampler.sample(root_nodes, ts)
        ret = sampler.get_ret()
        tot_time += ret[0].tot_time()
        ptr_time += ret[0].ptr_time()
        coo_time += ret[0].coo_time()
        sea_time += ret[0].search_time()
        sam_time += ret[0].sample_time()
        # for i in range(sample_config['history']):
        #     total_nodes += ret[i].dim_in() - ret[i].dim_out()
        #     unique_nodes += ret[i].dim_in() - ret[i].dim_out()
        #     if ret[i].dim_in() > ret[i].dim_out():
        #         ts = torch.from_numpy(ret[i].ts()[ret[i].dim_out():])
        #         nid = torch.from_numpy(ret[i].nodes()[ret[i].dim_out():]).float()
        #         nts = torch.stack([ts,nid],dim=1).cuda()
        #         uni_t_s = time.time()
        #         unts, idx = torch.unique(nts, dim=0, return_inverse=True)
        #         uni_time += time.time() - uni_t_s
        #         total_nodes += idx.shape[0]
        #         unique_nodes += unts.shape[0]

    print('total time  : {:.4f}'.format(tot_time))
    print('pointer time: {:.4f}'.format(ptr_time))
    print('coo time    : {:.4f}'.format(coo_time))
    print('search time : {:.4f}'.format(sea_time))
    print('sample time : {:.4f}'.format(sam_time))
    # print('unique time : {:.4f}'.format(uni_time))
    # print('unique per  : {:.4f}'.format(unique_nodes / total_nodes))


    
sampler.py
amazon-science/disttgl
from glob import glob
from setuptools import setup
from pybind11.setup_helpers import Pybind11Extension

ext_modules = [
    Pybind11Extension("sampler_core", 
                      ['sampler_core.cpp'],
                      extra_compile_args = ['-fopenmp'],
                      extra_link_args = ['-fopenmp'],),
]

setup(
    name = "sampler_core",
    version = "0.0.1",
    author = "Hongkuan Zhou",
    author_email = "hongkuaz@usc.edu",
    url = "https://tedzhouhk.github.io/about/",
    description = "Parallel Sampling for Temporal Graphs",
    ext_modules = ext_modules,
)
setup.py
amazon-science/disttgl
# Distributed Training

import argparse
import os

parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, help='dataset name')
parser.add_argument('--seed', type=int, default=0, help='random seed to use')
parser.add_argument('--omp_num_threads', type=int, default=6)
parser.add_argument('--train_neg_samples', type=int, default=1)
parser.add_argument('--batchsize', type=int, default=0)
parser.add_argument('--group', type=int, default=1, help='number of training group')
parser.add_argument('--minibatch_parallelism', type=int, default=1, help='how many mini-batch parallelism to use')
parser.add_argument('--neg_sets', type=int, default=32, help='how many number of negative sets in training group')
parser.add_argument('--edge_classification', action='store_true', help='whether to train with the edge classification task')
parser.add_argument('--edge_classes', type=int, default=0, help='how many class')
parser.add_argument('--log', type=str, default='', help='output log file')
parser.add_argument('--neg_rng', type=int, default=0, help='how many rngs to use in negative samplers')
parser.add_argument('--pbar', action='store_true', help='whether to add a progress bar')
parser.add_argument('--profile', action='store_true', help='whether to profile')
parser.add_argument('--partial_eval', action='store_true', help='whether to perform validation and test on partial dataset')
parser.add_argument('--partial_eval_interval', type=int, default=2000, help='how often to perform partial evaluation during training')
args = parser.parse_args()

if args.data in ['GDELT', 'LINK']:
    args.partial_eval = True
    args.neg_sets = 16
    args.pbar = True

if args.data in ['GDELT']:
    args.edge_classification = True
    args.edge_classes = 52

local_rank = int(os.environ['LOCAL_RANK'])
global_rank = int(os.environ['RANK'])
tot_rank = int(os.environ['WORLD_SIZE'])

os.environ['OMP_NUM_THREADS'] = str(args.omp_num_threads)
os.environ['MKL_NUM_THREADS'] = str(args.omp_num_threads)

# os.environ['TORCH_DISTRIBUTED_DEBUG'] = 'INFO'

import torch
import dgl
import datetime
import random
import pickle
import math
import threading
import numpy as np
from modules import *
from sampler import *
from utils import *
from get_config import *
from dataloader import *
from mailbox_daemon import *
from multiprocessing import Process
from tqdm import tqdm
from pathlib import Path
from contextlib import nullcontext
from sklearn.metrics import average_precision_score, roc_auc_score
from dgl.utils.shared_mem import create_shared_mem_array, get_shared_mem_array

torch.cuda.set_device(local_rank)

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed(args.seed)
torch.distributed.init_process_group(backend='gloo', timeout=datetime.timedelta(0, 3600 * 3))
nccl_group = torch.distributed.new_group(backend='nccl', timeout=datetime.timedelta(0, 3600 * 3))

if local_rank == 0:
    if args.log != '':
        log_dir = str(Path(args.log).parent)
        if not os.path.isdir(log_dir):
            os.makedirs(log_dir)
        log_f = open(args.log, 'w')

if global_rank == 0:
    if not os.path.isdir('models'):
        os.mkdir('models')
    path_saver = ['models/{}_{}.pkl'.format(args.data, time.time())]
    profile_path_saver = ['tb_log/{}_{}.pkl'.format(args.data, time.time())]
else:
    path_saver = [None]
    profile_path_saver = [None]
torch.distributed.broadcast_object_list(path_saver, src=0)
torch.distributed.broadcast_object_list(profile_path_saver, src=0)
path_saver = path_saver[0]
profile_path_saver = profile_path_saver[0]

if local_rank == 0:
    _node_feats, _edge_feats = load_feat(args.data)
dim_feats = [0, 0, 0, 0, 0, 0, 0, 0]
if local_rank == 0:
    if _node_feats is not None:
        dim_feats[0] = _node_feats.shape[0]
        dim_feats[1] = _node_feats.shape[1]
        dim_feats[2] = _node_feats.dtype
        dim_feats[6] = 1
        node_feats = create_shared_mem_array('node_feats', _node_feats.shape, dtype=_node_feats.dtype)
        node_feats.copy_(_node_feats)
        del _node_feats
    else:
        node_feats = None
    if _edge_feats is not None:
        dim_feats[3] = _edge_feats.shape[0]
        dim_feats[4] = _edge_feats.shape[1]
        dim_feats[5] = _edge_feats.dtype
        dim_feats[7] = 1
        edge_feats = create_shared_mem_array('edge_feats', _edge_feats.shape, dtype=_edge_feats.dtype)
        edge_feats.copy_(_edge_feats)
        del _edge_feats
    else: 
        edge_feats = None
torch.distributed.barrier()
torch.distributed.broadcast_object_list(dim_feats, src=0)
if local_rank > 0:
    node_feats = None
    edge_feats = None
    if dim_feats[6] == 1:
        node_feats = get_shared_mem_array('node_feats', (dim_feats[0], dim_feats[1]), dtype=dim_feats[2])
    if dim_feats[7] == 1:
        edge_feats = get_shared_mem_array('edge_feats', (dim_feats[3], dim_feats[4]), dtype=dim_feats[5])

sample_param, memory_param, gnn_param, train_param = get_config(args.data, tot_rank, minibatch_parallelism=args.minibatch_parallelism)
if args.train_neg_samples > 0:
    train_param['train_neg_samples'] = args.train_neg_samples
if args.batchsize > 0:
    train_param['lr'] = train_param['lr'] * args.batchsize / train_param['batch_size']
    train_param['batch_size'] = args.batchsize

with open('minibatches/{}_stats.pkl'.format(args.data), 'rb') as f:
    data_stats = pickle.load(f)
num_nodes = data_stats['num_nodes']

tot_group_rank = tot_rank // args.group
group_rank = global_rank % tot_group_rank
group_id = global_rank // tot_group_rank
# print('global_rank:{}, tot_group_rank:{}, group_id:{}'.format(global_rank,tot_group_rank,group_id))

# for minibatch parallelism
mb_tot_rank = tot_rank // args.minibatch_parallelism
mb_global_rank = global_rank // args.minibatch_parallelism
mb_local_rank = local_rank // args.minibatch_parallelism
mb_offset = global_rank % args.minibatch_parallelism
mb_tot_group_rank = mb_tot_rank // args.group
mb_group_rank = mb_global_rank % mb_tot_group_rank
mb_group_id = mb_global_rank // mb_tot_group_rank

mailbox = None
if memory_param['type'] != 'none':
    max_read_nodes = train_param['batch_size'] * (2 + train_param['train_neg_samples']) * (sample_param['neighbor'][0] + 1)
    max_write_nodes = train_param['batch_size'] * (2 + train_param['train_neg_samples'])
    if group_rank == 0:
        node_memory = create_shared_mem_array(str(group_id) + 'node_memory', torch.Size([num_nodes, memory_param['dim_out']]), dtype=torch.float32)
        mails = create_shared_mem_array(str(group_id) + 'mails', torch.Size([num_nodes, 2 * memory_param['dim_out']]), dtype=torch.float32)

        memory_read_buffer = create_shared_mem_array('{}_memory_read_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes, memory_param['dim_out']]), dtype=torch.float32)
        mail_read_buffer = create_shared_mem_array('{}_mail_read_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes, 2 * memory_param['dim_out']]), dtype=torch.float32)
        read_1idx_buffer = create_shared_mem_array('{}_read_1idx_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes + 1]), dtype=torch.int64)
        read_status = create_shared_mem_array('{}_read_status'.format(group_id), torch.Size([tot_group_rank]), dtype=torch.int)
        memory_write_buffer = create_shared_mem_array('{}_memory_write_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes, memory_param['dim_out']]), dtype=torch.float32)
        mail_write_buffer = create_shared_mem_array('{}_mail_write_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes, 2 * memory_param['dim_out']]), dtype=torch.float32)
        write_1idx_buffer = create_shared_mem_array('{}_write_1idx_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes + 1]), dtype=torch.int64)
        write_status = create_shared_mem_array('{}_write_status'.format(group_id), torch.Size([tot_group_rank]), dtype=torch.int)
        reset_status = create_shared_mem_array('{}_reset_status'.format(group_id), torch.Size([1]), dtype=torch.int)

        dims_mailbox = [num_nodes, memory_param['dim_out'], max_read_nodes, max_write_nodes, tot_group_rank, mb_tot_group_rank]

        read_status.zero_()
        write_status.zero_()
        reset_status.zero_()
        mailbox_daemon = Process(target=start_mailbox_daemon, args=(args.omp_num_threads, group_id, dims_mailbox, args.edge_classification))
        mailbox_daemon.start()
        torch.distributed.barrier()
    else:
        torch.distributed.barrier()
        node_memory = get_shared_mem_array(str(group_id) + 'node_memory', torch.Size([num_nodes, memory_param['dim_out']]), dtype=torch.float32)
        mails = get_shared_mem_array(str(group_id) + 'mails', torch.Size([num_nodes, 2 * memory_param['dim_out']]), dtype=torch.float32)

        memory_read_buffer = get_shared_mem_array('{}_memory_read_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes, memory_param['dim_out']]), dtype=torch.float32)
        mail_read_buffer = get_shared_mem_array('{}_mail_read_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes, 2 * memory_param['dim_out']]), dtype=torch.float32)
        read_1idx_buffer = get_shared_mem_array('{}_read_1idx_buffer'.format(group_id), torch.Size([tot_group_rank, mb_tot_group_rank, max_read_nodes + 1]), dtype=torch.int64)
        read_status = get_shared_mem_array('{}_read_status'.format(group_id), torch.Size([tot_group_rank]), dtype=torch.int)
        memory_write_buffer = get_shared_mem_array('{}_memory_write_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes, memory_param['dim_out']]), dtype=torch.float32)
        mail_write_buffer = get_shared_mem_array('{}_mail_write_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes, 2 * memory_param['dim_out']]), dtype=torch.float32)
        write_1idx_buffer = get_shared_mem_array('{}_write_1idx_buffer'.format(group_id), torch.Size([tot_group_rank, max_write_nodes + 1]), dtype=torch.int64)
        write_status = get_shared_mem_array('{}_write_status'.format(group_id), torch.Size([tot_group_rank]), dtype=torch.int)
        reset_status = get_shared_mem_array('{}_reset_status'.format(group_id), torch.Size([1]), dtype=torch.int)
    mailbox = MailBox(memory_param, num_nodes, dim_feats[4], node_memory, mails)
# print('here0: read_status:', read_status.data_ptr(), 'write_status', write_status.data_ptr())

if args.partial_eval:
    if local_rank == 0:
        mailbox_eval = MailBox(memory_param, num_nodes, dim_feats[4])
torch.distributed.barrier()

model = GeneralModel(dim_feats[1], dim_feats[4], sample_param, memory_param, gnn_param, train_param, num_nodes, edge_classification=args.edge_classification, edge_classes=args.edge_classes).cuda()
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True, process_group=nccl_group)

creterion = torch.nn.CrossEntropyLoss()

if memory_param['memory_update'] == 'smart' and not ('choice' in memory_param and memory_param['choice'] == 'deg'):
    optimizer = torch.optim.Adam(
        [
            {'params': [param for name, param in model.named_parameters() if param.requires_grad and name not in ['module.memory_updater.choice.weight', 'module.memory_updater.constant.weight']]},
        ],
        lr=train_param['lr'],
    )
else:
    raise TypeError

train_dataloader = DataLoader(args.data, train_param['train_neg_samples'], train_param['eval_neg_samples'], args.neg_sets, 'train', minibatch_parallelism=args.minibatch_parallelism, mailbox=mailbox, node_feats=node_feats, edge_feats=edge_feats, edge_classification=args.edge_classification)
val_dataloader = DataLoader(args.data, train_param['train_neg_samples'], train_param['eval_neg_samples'], args.neg_sets, 'val', edge_classification=args.edge_classification)
test_dataloader = DataLoader(args.data, train_param['train_neg_samples'], train_param['eval_neg_samples'], args.neg_sets, 'test', edge_classification=args.edge_classification)

def eval(mode='val', pos_only=False):
    model.eval()
    aps = list()
    mrrs = list()
    ec_pred = list()
    ec_true = list()
    f1mic = 0
    eval_neg_samples = train_param['eval_neg_samples']
    if mode == 'val':
        dataloader = val_dataloader
    elif mode == 'test':
        dataloader = test_dataloader
    elif mode == 'train':
        dataloader = train_dataloader
    if args.pbar:
        eval_pbar = tqdm(total=dataloader.tot_length, position=1)
    with torch.no_grad():
        total_loss = 0
        for pos_mfg, neg_mfg in dataloader:
            if pos_only or args.edge_classification:
                neg_mfg = None
                pos_mfg.combined = False
                mfg = pos_mfg
            else:
                mfg = combine_mfgs(pos_mfg, neg_mfg)
            mailbox.prep_input_mails(mfg)
            prepare_input(mfg, node_feats, edge_feats)
            mfg = mfg_to_cuda(mfg)

            # mailbox.prep_input_mails(pos_mfg)
            # mailbox.prep_input_mails(neg_mfg)
            # prepare_input(pos_mfg, node_feats, edge_feats)
            # prepare_input(neg_mfg, node_feats, edge_feats)
            # pos_mfg = mfg_to_cuda(pos_mfg)
            # neg_mfg = mfg_to_cuda(neg_mfg)
            # pred_pos, pred_neg = model(mfg, pos_mfg, neg_mfg)

            pred_pos, pred_neg = model(mfg)
            if not pos_only:
                if not args.edge_classification:
                    pred = torch.cat([pred_pos.squeeze(), pred_neg.squeeze()], dim=0).reshape(eval_neg_samples + 1, -1).T
                    loss = creterion(pred, torch.zeros(pred.shape[0], dtype=torch.long).cuda())
                    total_loss += loss * pred.shape[0]
                    y_pred = torch.cat([pred_pos, pred_neg], dim=0).sigmoid().cpu()
                    y_true = torch.cat([torch.ones(pred_pos.size(0)), torch.zeros(pred_neg.size(0))], dim=0)
                    aps.append(average_precision_score(y_true, y_pred))
                    mrrs.append(torch.reciprocal(torch.sum(pred_pos.squeeze() < pred_neg.squeeze().reshape(eval_neg_samples, -1), dim=0) + 1).type(torch.float))
                else:
                    pred = pred_pos
                    loss = creterion(pred, mfg.edge_cls_cuda)
                    total_loss += loss * pred.shape[0]
                    ec_pred.append(pred.cpu().numpy())
                    ec_true.append(mfg.edge_cls)

            mailbox.update_memory_and_mailbox(mfg, model.module.memory_updater.last_updated_memory)

            if args.pbar:
                eval_pbar.update(1)
    if not pos_only:
        if not args.edge_classification:
            ap = float(torch.tensor(aps).mean())
            mrr = float(torch.cat(mrrs).mean())
        else:
            ap = 0
            mrr = 0
            f1mic = calc_f1_mic(np.concatenate(ec_true), np.concatenate(ec_pred))
    else:
        ap = 0
        mrr = 0
    return total_loss, ap, mrr, f1mic

def partial_eval(mode='val', start_batch=3000, tot_length=500, prologue_length=100):
    model.eval()
    aps = list()
    mrrs = list()
    ec_pred = list()
    ec_true = list()
    f1mic = 0
    eval_neg_samples = train_param['eval_neg_samples']
    if mode == 'val':
        dataloader = val_dataloader
    elif mode == 'test':
        dataloader = test_dataloader
    mailbox_eval.reset()

    if args.pbar:
        eval_pbar = tqdm(total=tot_length, position=1)

    with torch.no_grad():
        total_loss = 0
        for i in range(start_batch, start_batch + tot_length):
            pos_mfg, neg_mfgs = dataloader.get(i)
            if not args.edge_classification:
                neg_mfg = neg_mfgs[0]
                mfg = combine_mfgs(pos_mfg, neg_mfg)
            else:
                mfg = pos_mfg
                mfg.combined = False
            mailbox_eval.prep_input_mails(mfg)
            prepare_input(mfg, node_feats, edge_feats)
            mfg = mfg_to_cuda(mfg)

            pred_pos, pred_neg = model(mfg)
            if i > start_batch + prologue_length:
                if not args.edge_classification:
                    pred = torch.cat([pred_pos.squeeze(), pred_neg.squeeze()], dim=0).reshape(eval_neg_samples + 1, -1).T
                    loss = creterion(pred, torch.zeros(pred.shape[0], dtype=torch.long).cuda())
                    total_loss += loss * pred.shape[0]
                    y_pred = torch.cat([pred_pos, pred_neg], dim=0).sigmoid().cpu()
                    y_true = torch.cat([torch.ones(pred_pos.size(0)), torch.zeros(pred_neg.size(0))], dim=0)
                    aps.append(average_precision_score(y_true, y_pred))
                    mrrs.append(torch.reciprocal(torch.sum(pred_pos.squeeze() < pred_neg.squeeze().reshape(eval_neg_samples, -1), dim=0) + 1).type(torch.float))
                else:
                    pred = pred_pos
                    loss = creterion(pred, mfg.edge_cls_cuda)
                    total_loss += loss * pred.shape[0]
                    ec_pred.append(pred.cpu().numpy())
                    ec_true.append(mfg.edge_cls)

            mailbox_eval.update_memory_and_mailbox(mfg, model.module.memory_updater.last_updated_memory)

            if args.pbar:
                eval_pbar.update(1)
    if not args.edge_classification:
        ap = float(torch.tensor(aps).mean())
        mrr = float(torch.cat(mrrs).mean())
    else:
        ap = 0
        mrr = 0
        f1mic = calc_f1_mic(np.concatenate(ec_true), np.concatenate(ec_pred))
    return total_loss, ap, mrr, f1mic

best_metric = 0
best_e = 0

train_iters = train_dataloader.tot_length + mb_tot_group_rank - 1
to_reset = False
just_reset = False
prefetch_init = True
last_it = -1
# ii = -tot_group_rank
# my_iter = 0
ii = -mb_tot_group_rank + train_iters // args.group * (args.group - mb_group_id)
my_iter = train_iters // args.group * (args.group - mb_group_id)
with torch.profiler.profile(
    schedule=torch.profiler.schedule(wait=10, warmup=10, active=16, skip_first=0, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler(profile_path_saver)
) if args.profile else nullcontext() as profiler:
    for e in range(train_param['epoch']):
        neg_samples = train_param['train_neg_samples']
        if local_rank == 0:
            print('Global Epoch: {:d}'.format(e))
            if args.log != '':
                log_f.write('Global Epoch: {:d}\n'.format(e))
            train_loss = 0
            if args.pbar:
                pbar = tqdm(total=train_iters, position=0)
        torch.distributed.barrier()
        # print(local_rank, 'here0')

        train_time = 0
        t_prep1 = 0
        t_prep2 = 0
        t_comput = 0
        t_writeback = 0
        for i in range(train_iters):
            if args.partial_eval and global_rank == 0:
                if i % args.partial_eval_interval == 0:
                    eval_loss, eval_ap, eval_mrr, eval_f1mic = partial_eval(mode='val')
                    if not args.edge_classification:
                        pbar.write('\tval loss:{:.4f}    val mrr:{:.4f}    val_ap:{:.4f}'.format(eval_loss, eval_mrr, eval_ap))
                    else:
                        pbar.write('\tval loss:{:.4f}    val f1mic:{:.4f}'.format(eval_loss, eval_f1mic))
                    if args.log != '':
                        if not args.edge_classification:
                            log_f.write('\tval loss:{:.4f}    val mrr:{:.4f}    val_ap:{:.4f}\n'.format(eval_loss, eval_mrr, eval_ap))
                        else:
                            log_f.write('\tval loss:{:.4f}    val f1mic:{:.4f}\n'.format(eval_loss, eval_f1mic))
                    curr_metric = eval_f1mic if args.edge_classification else eval_mrr
                    if curr_metric > best_metric:
                        best_e = e
                        best_metric = curr_metric
                        torch.save(model.state_dict(), path_saver)
                    torch.distributed.barrier()
            elif args.partial_eval:
                if i % args.partial_eval_interval == 0:
                    torch.distributed.barrier()

            # print('{}:iter {}/{} here0'.format(local_rank, i, train_iters - 1))

            model.train()
            if i == train_iters // args.group * mb_group_id:
                ii = -mb_tot_group_rank
                my_iter = 0
            if my_iter % mb_tot_group_rank == mb_group_rank:
                ii += mb_tot_group_rank
            it = ii + mb_group_rank
            if prefetch_init or (it <= 0 and to_reset):
                # print('{} init prefetch'.format(local_rank))
                train_dataloader.init_prefetch(it, num_neg=mb_tot_group_rank, offset=mb_offset, prefetch_interval=mb_tot_group_rank, rank=group_rank, memory_read_buffer=memory_read_buffer, mail_read_buffer=mail_read_buffer, read_1idx_buffer=read_1idx_buffer, read_status=read_status)
                prefetch_init = False
                last_it = -1
            it = 0 if it < 0 else it 
            it = train_dataloader.tot_length - 1 if it >= train_dataloader.tot_length else it

            # for r in range(tot_rank):
            #     if local_rank == r:
            #         print('{}:{} iter'.format(local_rank, it))
            #     torch.distributed.barrier()

            # print('{}:iter {}/{} here2'.format(local_rank, i, train_iters - 1))
            if just_reset:
                just_reset = False
                reset_status[0] = 0
            if it == 0 and to_reset:
                if mb_group_rank == 0 and mb_offset == 0:
                    # print('{} set reset to -1'.format(local_rank))
                    reset_status[0] = 1
                to_reset = False
                just_reset = True
                # read_status[group_rank] = -1
            elif not to_reset:
                if it > 0:
                    to_reset = True
            if just_reset:
                # a virtual barrier for all reset trainers
                while reset_status[0] != 2:
                    pass
            # torch.distributed.barrier()

            # print('{}:iter {}/{} here1'.format(local_rank, i, train_iters - 1))
            t_s = time.time()
            t_train_s = time.time()
            if last_it != it:
                train_dataloader.prefetch_next()
                # torch.distributed.barrier()
                t_prep1 += time.time() - t_s
                t_s = time.time()
                last_it = it
                if args.edge_classification:
                    mfg = train_dataloader.get_fetched()
                    # mailbox.prep_zero_mails(mfg)
                    mfg = mfg_to_cuda(mfg)
            else:
                # torch.distributed.barrier()
                t_prep1 += time.time() - t_s
                t_s = time.time()
            
            
            if not args.edge_classification:
                mfg = train_dataloader.get_fetched()
                # mailbox.prep_zero_mails(mfg)
                mfg = mfg_to_cuda(mfg)
            # print('{} mfg ready at iter {}'.format(local_rank, i))
            # torch.distributed.barrier()
            t_prep2 += time.time() - t_s
            
            t_s = time.time()
            optimizer.zero_grad()
            # torch.cuda.default_stream().wait_stream(mfg.stream)
            if my_iter % mb_tot_group_rank == mb_group_rank:
                write_buffer = WriteBuffer(group_rank, memory_write_buffer, mail_write_buffer, write_1idx_buffer, write_status)
                # print('{} set write_buffer[{}] at iter {}'.format(local_rank, group_rank, i))
                pred_pos, pred_neg = model(mfg, write_buffer=write_buffer)
            else:
                pred_pos, pred_neg = model(mfg)
            if not args.edge_classification:
                pred = torch.cat([pred_pos.squeeze(), pred_neg.squeeze()], dim=0).reshape(neg_samples + 1, -1).T
                loss = creterion(pred, torch.zeros(pred.shape[0], dtype=torch.long, device=pred.device))
            else:
                pred = pred_pos
                loss = creterion(pred, mfg.edge_cls_cuda)
            loss.backward()
            optimizer.step()
            if args.profile:
                profiler.step()
            if local_rank == 0:
                train_loss += float(loss) * train_param['batch_size']
            # torch.distributed.barrier()
            # print('\tt_comput: {}ms'.format((time.time() - t_s) * 1000))
            t_comput += time.time() - t_s
            
            t_s = time.time()
            # if my_iter % mb_tot_group_rank == mb_group_rank:
            #     mailbox.update_memory_and_mailbox(mfg, model.module.memory_updater.last_updated_memory)
            # torch.distributed.barrier()
            t_writeback += time.time() - t_s

            train_time += time.time() - t_train_s
            if local_rank == 0:
                if args.pbar:
                    pbar.update(1)

            my_iter += 1
        if local_rank == 0:
            print('\ttrain loss:{:.4f}  train time:{:.2f}'.format(train_loss, train_time))
            print('\tt_prep1:{:.2f} t_prep2:{:.2f} t_compute:{:.2f} t_writeback:{:.2f}'.format(t_prep1, t_prep2, t_comput, t_writeback))
            if args.log != '':
                log_f.write('\ttrain loss:{:.4f}  train time:{:.2f}\n'.format(train_loss, train_time))
                log_f.write('\tt_prep1:{:.2f} t_prep2:{:.2f} t_compute:{:.2f} t_writeback:{:.2f}\n'.format(t_prep1, t_prep2, t_comput, t_writeback))
        if not args.partial_eval:
            if global_rank == 0:
                eval_loss, eval_ap, eval_mrr, eval_f1mic = eval('val')
                if not args.edge_classification:
                    print('\tval loss:{:.4f}    val mrr:{:.4f}    val_ap:{:.4f}'.format(eval_loss, eval_mrr, eval_ap))
                else:
                    print('\tval loss:{:.4f}    val f1mic:{:.4f}'.format(eval_loss, eval_f1mic))
                if args.log != '':
                    if not args.edge_classification:
                        log_f.write('\tval loss:{:.4f}    val mrr:{:.4f}    val_ap:{:.4f}\n'.format(eval_loss, eval_mrr, eval_ap))
                    else:
                        log_f.write('\tval loss:{:.4f}    val f1mic:{:.4f}\n'.format(eval_loss, eval_f1mic))
                curr_metric = eval_f1mic if args.edge_classification else eval_mrr
                if curr_metric > best_metric:
                    best_e = e
                    best_metric = curr_metric
                    torch.save(model.state_dict(), path_saver)
                    _, test_ap, test_mrr, test_f1mic = eval('test')
                    print('\ttest AP:{:4f}  test MRR:{:4f}'.format(test_ap, test_mrr))
    
if group_rank == 0:
    mailbox_daemon.kill()

if global_rank == 0:
    print('Loading model at epoch {}...'.format(best_e))
    model.load_state_dict(torch.load(path_saver))
    model.eval()
    if args.partial_eval:
        _, test_ap, test_mrr, test_f1mic = partial_eval('test')
    if args.edge_classification:
        print('\ttest f1mic:{:4f}'.format(test_f1mic))
        if args.log != '':
            log_f.write('\ttest f1mic:{:4f}\n'.format(test_f1mic))
            log_f.close()
    else:
        print('\ttest AP:{:4f}  test MRR:{:4f}'.format(test_ap, test_mrr))
        if args.log != '':
            log_f.write('\ttest AP:{:4f}  test MRR:{:4f}\n'.format(test_ap, test_mrr))
            log_f.close()
torch.distributed.barrier()
train.py
amazon-science/disttgl
import torch
import os
import yaml
import dgl
import time
import pandas as pd
import numpy as np
from sklearn import metrics

class WriteBuffer:
    # this class is to avoid the shallow copy for the input parameters in the .forward() function in torch.distributed.model
    def __init__(self, rank, memory_write_buffer, mail_write_buffer, write_1idx_buffer, write_status):
        self.rank = rank
        self.memory_write_buffer = memory_write_buffer
        self.mail_write_buffer = mail_write_buffer
        self.write_1idx_buffer = write_1idx_buffer
        self.write_status = write_status

def calc_f1_mic(y_true, y_pred):
    y_pred[y_pred > 0] = 1
    y_pred[y_pred <= 0] = 0
    return metrics.f1_score(y_true, y_pred, average="micro")

def load_feat(d):
    node_feats = None
    if os.path.exists('DATA/{}/node_features.pt'.format(d)):
        node_feats = torch.load('DATA/{}/node_features.pt'.format(d))
    elif os.path.exists('DATA/{}/learned_node_feats.pt'.format(d)):
        node_feats = torch.load('DATA/{}/learned_node_feats.pt'.format(d))
    edge_feats = None
    if os.path.exists('DATA/{}/edge_features_e0.pt'.format(d)):
        edge_feats = torch.load('DATA/{}/edge_features_e0.pt'.format(d))
    elif os.path.exists('DATA/{}/edge_features.pt'.format(d)):
        edge_feats = torch.load('DATA/{}/edge_features.pt'.format(d))
        edge_feats = torch.cat([edge_feats, torch.zeros((1, edge_feats.shape[1]), dtype=edge_feats.dtype)])
        torch.save(edge_feats, 'DATA/{}/edge_features_e0.pt'.format(d))
    if node_feats is not None:
        if node_feats.dtype == torch.bool:
            node_feats = node_feats.type(torch.int8)
    if edge_feats is not None:
        if edge_feats.dtype == torch.bool:
            edge_feats = edge_feats.type(torch.int8)
    return node_feats, edge_feats

def load_graph(d):
    if os.path.isfile('DATA/{}/edges+uniq.csv'.format(d)):
        df = pd.read_csv('DATA/{}/edges+uniq.csv'.format(d))
    else:
        df = pd.read_csv('DATA/{}/edges.csv'.format(d))
    g = np.load('DATA/{}/ext_full.npz'.format(d))
    return g, df

def parse_config(f):
    conf = yaml.safe_load(open(f, 'r'))
    sample_param = conf['sampling'][0]
    memory_param = conf['memory'][0]
    gnn_param = conf['gnn'][0]
    train_param = conf['train'][0]
    return sample_param, memory_param, gnn_param, train_param

def unique_last_idx(a):
    # find index of last appearance of an element in a array
    uni, inv = torch.unique(a, return_inverse=True)
    perm = torch.arange(inv.size(0), dtype=inv.dtype, device=inv.device)
    perm = inv.new_empty(uni.size(0)).scatter_(0, inv, perm)
    return perm

def to_dgl_blocks(ret, hist, num_neg_samples, cuda=True, combine=True):
    mfgs = list()
    for r in ret:
        positive_nodes = r.dim_out() // (num_neg_samples + 2) * 2
        src_id = torch.from_numpy(r.nodes())
        src_uniq, src_uniq_idx = torch.unique(src_id, sorted=False, return_inverse=True)
        col = src_uniq_idx[torch.from_numpy(r.col()).long()]
        row = torch.from_numpy(r.row()).long()

        b = dgl.create_block((col, row), num_src_nodes=src_uniq.shape[0], num_dst_nodes=r.dim_out())
        b.srcdata['ID'] = src_uniq

        b.edata['dt'] = torch.from_numpy(r.dts())[b.num_dst_nodes():]
        b.edata['ID'] = torch.from_numpy(r.eid())

        b.positive_memory_idx_end = positive_nodes
        b.src_idx = src_uniq_idx

        if cuda:
            b.src_idx = b.src_idx.cuda()
            mfgs.append(b.to('cuda:0'))
        else:
            mfgs.append(b)
    mfgs = list(map(list, zip(*[iter(mfgs)] * hist)))
    mfgs.reverse()
    return mfgs

def combine_mfgs(pos_mfg, neg_mfg):
    col = torch.cat([pos_mfg.edges()[0], neg_mfg.edges()[0] + pos_mfg.num_src_nodes()])
    row = torch.cat([pos_mfg.edges()[1], neg_mfg.edges()[1] + pos_mfg.num_dst_nodes()])
    mfg = dgl.create_block((col, row), num_src_nodes=pos_mfg.num_src_nodes() + neg_mfg.num_src_nodes(), num_dst_nodes=pos_mfg.num_dst_nodes() + neg_mfg.num_dst_nodes())
    mfg.combined = True
    mfg.num_pos_src = pos_mfg.num_src_nodes()
    mfg.num_pos_dst = pos_mfg.num_dst_nodes()
    mfg.num_neg_src = neg_mfg.num_src_nodes()
    mfg.num_neg_dst = neg_mfg.num_dst_nodes()
    mfg.num_pos_idx = pos_mfg.src_idx.shape[0]
    mfg.num_neg_idx = neg_mfg.src_idx.shape[0]
    mfg.node_memory_mask = pos_mfg.node_memory_mask
    mfg.srcdata['ID'] = torch.cat([pos_mfg.srcdata['ID'], neg_mfg.srcdata['ID']])
    mfg.srcdata['mem_ts'] = torch.cat([pos_mfg.srcdata['mem_ts'], neg_mfg.srcdata['mem_ts']])
    mfg.srcdata['mail_ts'] = torch.cat([pos_mfg.srcdata['mail_ts'], neg_mfg.srcdata['mail_ts']])
    mfg.srcdata['mail_e'] = torch.cat([pos_mfg.srcdata['mail_e'], neg_mfg.srcdata['mail_e']])
    mfg.edata['dt'] = torch.cat([pos_mfg.edata['dt'], neg_mfg.edata['dt']])
    mfg.edata['ID'] = torch.cat([pos_mfg.edata['ID'], neg_mfg.edata['ID']])
    mfg.src_idx = torch.cat([pos_mfg.src_idx, neg_mfg.src_idx + pos_mfg.num_src_nodes()])
    return mfg

def node_to_dgl_blocks(root_nodes, ts, cuda=True):
    mfgs = list()
    b = dgl.create_block(([],[]), num_src_nodes=root_nodes.shape[0], num_dst_nodes=root_nodes.shape[0])
    b.srcdata['ID'] = torch.from_numpy(root_nodes)
    b.srcdata['ts'] = torch.from_numpy(ts)
    if cuda:
        mfgs.insert(0, [b.to('cuda:0')])
    else:
        mfgs.insert(0, [b])
    return mfgs

def mfg_to_cuda(mfg, cuda_device_id=None):
    if cuda_device_id is not None:
        mfg = mfg.to('cuda:{}'.format(cuda_device_id))
    else:
        mfg = mfg.to('cuda:{}'.format(torch.cuda.current_device()))
    mfg.src_idx_cuda = mfg.src_idx.cuda()
    if hasattr(mfg, 'edge_cls'):
        mfg.edge_cls_cuda = mfg.edge_cls.cuda()
    if hasattr(mfg, 'node_memory_mask'):
        mfg.node_memory_mask_cuda = mfg.node_memory_mask.cuda()
    return mfg

def prepare_input(mfg, node_feats, edge_feats):
    with torch.no_grad():
        if node_feats is not None:
            mfg.srcdata['rh'] = torch.index_select(node_feats, 0, mfg.srcdata['ID']).float()
        if edge_feats is not None:
            mfg.edata['f'] = torch.index_select(edge_feats, 0, mfg.edata['ID']).float()
            mfg.srcdata['mail_ef'] = torch.index_select(edge_feats, 0, mfg.srcdata['mail_e']).float()
    return

def get_ids(mfgs, node_feats, edge_feats):
    nids = list()
    eids = list()
    if node_feats is not None:
        for b in mfgs[0]:
            nids.append(b.srcdata['ID'].long())
    if edge_feats is not None:
        for mfg in mfgs:
            for b in mfg:
                eids.append(b.edata['ID'].long())
    return nids, eids

def get_pinned_buffers(sample_param, batch_size, node_feats, edge_feats):
    pinned_nfeat_buffs = list()
    pinned_efeat_buffs = list()
    limit = int(batch_size * 3.3)
    if 'neighbor' in sample_param:
        for i in sample_param['neighbor']:
            limit *= i + 1
            if edge_feats is not None:
                for _ in range(sample_param['history']):
                    pinned_efeat_buffs.insert(0, torch.zeros((limit, edge_feats.shape[1]), pin_memory=True))
    if node_feats is not None:
        for _ in range(sample_param['history']):
            pinned_nfeat_buffs.insert(0, torch.zeros((limit, node_feats.shape[1]), pin_memory=True))
    return pinned_nfeat_buffs, pinned_efeat_buffs

def send_recv_mfg_GPU(mfgs, src, dst, nccl_group):
    local_rank = int(os.environ['LOCAL_RANK'])
    # print(local_rank, src, dst)
    if local_rank == src:
        block = mfgs[0][0]
        dim_block = torch.tensor([block.num_src_nodes(), block.num_dst_nodes(), block.num_edges(), block.srcdata['mem'].shape[1], block.srcdata['mem_input'].shape[1], block.edata['f'].shape[1], block.src_idx.shape[0]])
        torch.distributed.send(dim_block, dst)
        torch.distributed.send(block.edges()[0], dst, group=nccl_group)
        torch.distributed.send(block.edges()[1], dst, group=nccl_group)

        torch.distributed.send(block.srcdata['ID'], dst, group=nccl_group)
        # torch.distributed.send(block.srcdata['ts'], dst, group=nccl_group)
        torch.distributed.send(block.srcdata['mem'], dst, group=nccl_group)
        torch.distributed.send(block.srcdata['mem_ts'], dst, group=nccl_group)
        torch.distributed.send(block.srcdata['mem_input'], dst, group=nccl_group)
        torch.distributed.send(block.srcdata['mail_ts'], dst, group=nccl_group)

        torch.distributed.send(block.edata['ID'], dst, group=nccl_group)
        torch.distributed.send(block.edata['dt'], dst, group=nccl_group)
        torch.distributed.send(block.edata['f'], dst, group=nccl_group)

        torch.distributed.send(block.src_idx, dst, group=nccl_group)

    elif local_rank == dst:
        dim_block = torch.tensor([0, 0, 0, 0, 0, 0, 0])
        torch.distributed.recv(dim_block, src)
        dim_block = dim_block.tolist()
        edges0 = torch.zeros(dim_block[2], dtype=torch.int64, device=torch.device('cuda:0'))
        edges1 = torch.zeros(dim_block[2], dtype=torch.int64, device=torch.device('cuda:0'))
        torch.distributed.recv(edges0, src, group=nccl_group)
        torch.distributed.recv(edges1, src, group=nccl_group)
        block = dgl.create_block((edges0, edges1), num_src_nodes=dim_block[0], num_dst_nodes=dim_block[1])

        src_id = torch.zeros(dim_block[0], dtype=torch.int32, device=torch.device('cuda:0'))
        # src_ts = torch.zeros(dim_block[0], dtype=torch.float32, device=torch.device('cuda:0'))
        src_mem = torch.zeros((dim_block[0], dim_block[3]), dtype=torch.float32, device=torch.device('cuda:0'))
        src_mem_ts = torch.zeros(dim_block[0], dtype=torch.float32, device=torch.device('cuda:0'))
        src_mem_input = torch.zeros((dim_block[0], dim_block[4]), dtype=torch.float32, device=torch.device('cuda:0'))
        src_mail_ts = torch.zeros((dim_block[0], 1), dtype=torch.float32, device=torch.device('cuda:0'))
        torch.distributed.recv(src_id, src, group=nccl_group)
        # torch.distributed.recv(src_ts, src, group=nccl_group)
        torch.distributed.recv(src_mem, src, group=nccl_group)
        torch.distributed.recv(src_mem_ts, src, group=nccl_group)
        torch.distributed.recv(src_mem_input, src, group=nccl_group)
        torch.distributed.recv(src_mail_ts, src, group=nccl_group)
        block.srcdata['ID'] = src_id
        # block.srcdata['ts'] = src_ts
        block.srcdata['mem'] = src_mem
        block.srcdata['mem_ts'] = src_mem_ts
        block.srcdata['mem_input'] = src_mem_input
        block.srcdata['mail_ts'] = src_mail_ts

        e_id = torch.zeros(dim_block[2], dtype=torch.int32, device=torch.device('cuda:0'))
        e_dt = torch.zeros(dim_block[2], dtype=torch.float32, device=torch.device('cuda:0'))
        e_f = torch.zeros((dim_block[2], dim_block[5]), dtype=torch.float32, device=torch.device('cuda:0'))
        torch.distributed.recv(e_id, src, group=nccl_group)
        torch.distributed.recv(e_dt, src, group=nccl_group)
        torch.distributed.recv(e_f, src, group=nccl_group)
        block.edata['ID'] = e_id
        block.edata['dt'] = e_dt
        block.edata['f'] = e_f

        src_idx = torch.zeros(dim_block[6], dtype=torch.int64, device=torch.device('cuda:0'))
        torch.distributed.recv(src_idx, src, group=nccl_group)
        block.src_idx = src_idx

    return [[block]]
utils.py
amazon-science/disttgl