Skip to content

S
starrygl-DynamicHistory
Commit 54424543 by zlj
Options

change the torch_utilsdependece path

parent 8768fe6f
Showing with 295 additions and 154 deletions
CMakeLists.txt
set(CMAKE_SOURCE_DIR "/mnt/data/")
cmake_minimum_required(VERSION 3.15) cmake_minimum_required(VERSION 3.15)
project(starrygl_ops VERSION 0.1) project(starrygl_ops VERSION 0.1)
...@@ -45,7 +46,7 @@ if(WITH_METIS) ...@@ -45,7 +46,7 @@ if(WITH_METIS)
add_definitions(-DWITH_METIS) add_definitions(-DWITH_METIS)
set(GKLIB_DIR "${CMAKE_SOURCE_DIR}/third_party/GKlib") set(GKLIB_DIR "${CMAKE_SOURCE_DIR}/third_party/GKlib")
set(METIS_DIR "${CMAKE_SOURCE_DIR}/third_party/METIS") set(METIS_DIR "${CMAKE_SOURCE_DIR}/third_party/METIS")
message(${METIS_DIR})
set(GKLIB_INCLUDE_DIRS "${GKLIB_DIR}/include") set(GKLIB_INCLUDE_DIRS "${GKLIB_DIR}/include")
file(GLOB_RECURSE GKLIB_LIBRARIES "${GKLIB_DIR}/lib/lib*.a") file(GLOB_RECURSE GKLIB_LIBRARIES "${GKLIB_DIR}/lib/lib*.a")
...@@ -104,10 +105,10 @@ endif() ...@@ -104,10 +105,10 @@ endif()
# add libsampler.so # add libsampler.so
set(SAMLPER_NAME "${PROJECT_NAME}_sampler") set(SAMLPER_NAME "${PROJECT_NAME}_sampler")
set(BOOST_INCLUDE_DIRS "${CMAKE_SOURCE_DIR}/third_party/boost_1_83_0")
include_directories(${BOOST_INCLUDE_DIRS})
file(GLOB_RECURSE SAMPLER_SRCS "csrc/sampler/*.cpp") file(GLOB_RECURSE SAMPLER_SRCS "csrc/sampler/*.cpp")
add_library(${SAMLPER_NAME} SHARED ${SAMPLER_SRCS}) add_library(${SAMLPER_NAME} SHARED ${SAMPLER_SRCS})
target_include_directories(${SAMLPER_NAME} PRIVATE "csrc/sampler/include") target_include_directories(${SAMLPER_NAME} PRIVATE "csrc/sampler/include")
target_compile_options(${SAMLPER_NAME} PRIVATE -O3) target_compile_options(${SAMLPER_NAME} PRIVATE -O3)
......
Makefile 0 → 100644
This diff is collapsed. Click to expand it.
config/TGN.yml
...@@ -23,7 +23,7 @@ gnn: ...@@ -23,7 +23,7 @@ gnn:
dim_time: 100 dim_time: 100
dim_out: 100 dim_out: 100
train: train:
- epoch: 10 - epoch: 5
#batch_size: 100 #batch_size: 100
# reorder: 16 # reorder: 16
lr: 0.0001 lr: 0.0001
......
csrc/sampler/export.cpp
...@@ -2,6 +2,7 @@ ...@@ -2,6 +2,7 @@
#include <sampler.h> #include <sampler.h>
#include <output.h> #include <output.h>
#include <neighbors.h> #include <neighbors.h>
#include <temporal_utils.h>
/*------------Python Bind--------------------------------------------------------------*/ /*------------Python Bind--------------------------------------------------------------*/
...@@ -16,7 +17,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) ...@@ -16,7 +17,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
py::return_value_policy::reference) py::return_value_policy::reference)
.def("divide_nodes_to_part", .def("divide_nodes_to_part",
&divide_nodes_to_part, &divide_nodes_to_part,
py::return_value_policy::reference); py::return_value_policy::reference)
.def("sparse_get_index",
&sparse_get_index,
py::return_value_policy::reference)
.def("get_norm_temporal",
&get_norm_temporal,
py::return_value_policy::reference
);
py::class_<TemporalGraphBlock>(m, "TemporalGraphBlock") py::class_<TemporalGraphBlock>(m, "TemporalGraphBlock")
.def(py::init<vector<NodeIDType> &, vector<NodeIDType> &, .def(py::init<vector<NodeIDType> &, vector<NodeIDType> &,
...@@ -79,4 +87,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) ...@@ -79,4 +87,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
.def("neighbor_sample_from_nodes", &ParallelSampler::neighbor_sample_from_nodes) .def("neighbor_sample_from_nodes", &ParallelSampler::neighbor_sample_from_nodes)
.def("reset", &ParallelSampler::reset) .def("reset", &ParallelSampler::reset)
.def("get_ret", [](const ParallelSampler &ps) { return ps.ret; }); .def("get_ret", [](const ParallelSampler &ps) { return ps.ret; });
} }
\ No newline at end of file
starrygl/sample/part_utils/torch_utils.cpp → csrc/sampler/include/temporal_utils.h
#include <pybind11/pybind11.h> #pragma once
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include <torch/extension.h> #include <torch/extension.h>
#include <parallel_hashmap/phmap.h> #include <parallel_hashmap/phmap.h>
#include <cstring> #include <cstring>
......
data_maker.py
...@@ -44,14 +44,14 @@ def load_feat(d, rand_de=0, rand_dn=0): ...@@ -44,14 +44,14 @@ def load_feat(d, rand_de=0, rand_dn=0):
data_name = args.data_name data_name = args.data_name
g = np.load('../tgl_main/DATA/'+data_name+'/ext_full.npz') g = np.load('/mnt/nfs/fzz/TGL-DATA/'+data_name+'/ext_full.npz')
df = pd.read_csv('../tgl_main/DATA/'+data_name+'/edges.csv') df = pd.read_csv('/mnt/nfs/fzz/TGL-DATA/'+data_name+'/edges.csv')
if os.path.exists('../tgl_main/DATA/'+data_name+'/node_features.pt'): if os.path.exists('/mnt/nfs/fzz/TGL-DATA/'+data_name+'/node_features.pt'):
n_feat = torch.load('../tgl_main/DATA/'+data_name+'/node_features.pt') n_feat = torch.load('/mnt/nfs/fzz/TGL-DATA/'+data_name+'/node_features.pt')
else: else:
n_feat = None n_feat = None
if os.path.exists('../tgl_main/DATA/'+data_name+'/edge_features.pt'): if os.path.exists('/mnt/nfs/fzz/TGL-DATA/'+data_name+'/edge_features.pt'):
e_feat = torch.load('../tgl_main/DATA/'+data_name+'/edge_features.pt') e_feat = torch.load('/mnt/nfs/fzz/TGL-DATA/'+data_name+'/edge_features.pt')
else: else:
e_feat = None e_feat = None
......
detect_cuda_compute_capabilities.cu 0 → 100644
#include <cuda_runtime.h>
#include <cstdio>
int main()
{
int count = 0;
if (cudaSuccess != cudaGetDeviceCount(&count)) return -1;
if (count == 0) return -1;
for (int device = 0; device < count; ++device)
{
cudaDeviceProp prop;
if (cudaSuccess == cudaGetDeviceProperties(&prop, device))
std::printf("%d.%d ", prop.major, prop.minor);
}
return 0;
}
detect_cuda_version.cc 0 → 100644
#include <cuda.h>
#include <cstdio>
int main() {
printf("%d.%d", CUDA_VERSION / 1000, (CUDA_VERSION / 10) % 100);
return 0;
}
install.sh
...@@ -3,9 +3,9 @@ ...@@ -3,9 +3,9 @@
mkdir -p build && cd build mkdir -p build && cd build
cmake .. \ cmake .. \
-DCMAKE_EXPORT_COMPILE_COMMANDS=ON \ -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
-DCMAKE_PREFIX_PATH="/home/hwj/.miniconda3/envs/sgl/lib/python3.10/site-packages" \ -DCMAKE_PREFIX_PATH="/home/zlj/.miniconda3/envs/dgnn/lib/python3.10/site-packages" \
-DPython3_ROOT_DIR="/home/hwj/.miniconda3/envs/sgl" \ -DPython3_ROOT_DIR="/home/zlj/.miniconda3/envs/dgnn" \
-DCUDA_TOOLKIT_ROOT_DIR="/home/hwj/.local/cuda-11.8" \ -DCUDA_TOOLKIT_ROOT_DIR="/home/zlj/local/cuda-12.2" \
&& make -j32 \ && make -j32 \
&& rm -rf ../starrygl/lib \ && rm -rf ../starrygl/lib \
&& mkdir ../starrygl/lib \ && mkdir ../starrygl/lib \
......
starrygl/distributed/cclib.py
...@@ -133,7 +133,6 @@ def all_to_all_s( ...@@ -133,7 +133,6 @@ def all_to_all_s(
output_sizes = [t-s for s, t in zip(output_rowptr, output_rowptr[1:])] output_sizes = [t-s for s, t in zip(output_rowptr, output_rowptr[1:])]
input_sizes = [t-s for s, t in zip(input_rowptr, input_rowptr[1:])] input_sizes = [t-s for s, t in zip(input_rowptr, input_rowptr[1:])]
return dist.all_to_all_single( return dist.all_to_all_single(
output=output_tensor, output=output_tensor,
input=input_tensor, input=input_tensor,
......
starrygl/distributed/utils.py
...@@ -17,6 +17,7 @@ class TensorAccessor: ...@@ -17,6 +17,7 @@ class TensorAccessor:
self._data = data self._data = data
self._ctx = DistributedContext.get_default_context() self._ctx = DistributedContext.get_default_context()
self._rref = rpc.RRef(data) self._rref = rpc.RRef(data)
self.stream = torch.cuda.Stream()
@property @property
def data(self): def data(self):
...@@ -29,6 +30,12 @@ class TensorAccessor: ...@@ -29,6 +30,12 @@ class TensorAccessor:
@property @property
def ctx(self): def ctx(self):
return self._ctx return self._ctx
@staticmethod
@rpc.functions.async_execution
def _index_selet(self):
fut = torch.futures.Future()
fut.set_result(None)
return fut
def all_gather_rrefs(self) -> List[rpc.RRef]: def all_gather_rrefs(self) -> List[rpc.RRef]:
return self.ctx.all_gather_remote_objects(self.rref) return self.ctx.all_gather_remote_objects(self.rref)
...@@ -180,7 +187,7 @@ class DistributedTensor: ...@@ -180,7 +187,7 @@ class DistributedTensor:
def ctx(self): def ctx(self):
return self.accessor.ctx return self.accessor.ctx
def all_to_all_ind2ptr(self, dist_index: Union[Tensor, DistIndex]) -> Dict[str, Union[List[int], Tensor]]: def all_to_all_ind2ptr(self, dist_index: Union[Tensor, DistIndex],group = None) -> Dict[str, Union[List[int], Tensor]]:
if isinstance(dist_index, Tensor): if isinstance(dist_index, Tensor):
dist_index = DistIndex(dist_index) dist_index = DistIndex(dist_index)
send_ptr = torch.ops.torch_sparse.ind2ptr(dist_index.part, self.num_parts()) send_ptr = torch.ops.torch_sparse.ind2ptr(dist_index.part, self.num_parts())
...@@ -196,7 +203,7 @@ class DistributedTensor: ...@@ -196,7 +203,7 @@ class DistributedTensor:
recv_ptr = recv_ptr.tolist() recv_ptr = recv_ptr.tolist()
recv_ind = torch.full((recv_ptr[-1],), (2**62-1)*2+1, dtype=dist_index.dtype, device=dist_index.device) recv_ind = torch.full((recv_ptr[-1],), (2**62-1)*2+1, dtype=dist_index.dtype, device=dist_index.device)
all_to_all_s(recv_ind, dist_index.loc, send_ptr, recv_ptr) all_to_all_s(recv_ind, dist_index.loc, recv_ptr, send_ptr,group=group)
return { return {
"send_ptr": send_ptr, "send_ptr": send_ptr,
...@@ -209,6 +216,7 @@ class DistributedTensor: ...@@ -209,6 +216,7 @@ class DistributedTensor:
send_ptr: Optional[List[int]] = None, send_ptr: Optional[List[int]] = None,
recv_ptr: Optional[List[int]] = None, recv_ptr: Optional[List[int]] = None,
recv_ind: Optional[List[int]] = None, recv_ind: Optional[List[int]] = None,
group = None
) -> Tensor: ) -> Tensor:
if dist_index is not None: if dist_index is not None:
dist_dict = self.all_to_all_ind2ptr(dist_index) dist_dict = self.all_to_all_ind2ptr(dist_index)
...@@ -218,7 +226,7 @@ class DistributedTensor: ...@@ -218,7 +226,7 @@ class DistributedTensor:
data = self.accessor.data[recv_ind] data = self.accessor.data[recv_ind]
recv = torch.empty(send_ptr[-1], *data.shape[1:], dtype=data.dtype, device=data.device) recv = torch.empty(send_ptr[-1], *data.shape[1:], dtype=data.dtype, device=data.device)
all_to_all_s(recv, data, send_ptr, recv_ptr) all_to_all_s(recv, data, send_ptr, recv_ptr,group=group)
return recv return recv
def all_to_all_set(self, def all_to_all_set(self,
...@@ -227,6 +235,7 @@ class DistributedTensor: ...@@ -227,6 +235,7 @@ class DistributedTensor:
send_ptr: Optional[List[int]] = None, send_ptr: Optional[List[int]] = None,
recv_ptr: Optional[List[int]] = None, recv_ptr: Optional[List[int]] = None,
recv_ind: Optional[List[int]] = None, recv_ind: Optional[List[int]] = None,
group = None
): ):
if dist_index is not None: if dist_index is not None:
dist_dict = self.all_to_all_ind2ptr(dist_index) dist_dict = self.all_to_all_ind2ptr(dist_index)
...@@ -235,7 +244,7 @@ class DistributedTensor: ...@@ -235,7 +244,7 @@ class DistributedTensor:
recv_ind = dist_dict["recv_ind"] recv_ind = dist_dict["recv_ind"]
recv = torch.empty(recv_ptr[-1], *data.shape[1:], dtype=data.dtype, device=data.device) recv = torch.empty(recv_ptr[-1], *data.shape[1:], dtype=data.dtype, device=data.device)
all_to_all_s(recv, data, recv_ptr, send_ptr) all_to_all_s(recv, data, recv_ptr, send_ptr,group=group)
self.accessor.data.index_copy_(0, recv_ind, recv) self.accessor.data.index_copy_(0, recv_ind, recv)
def index_select(self, dist_index: Union[Tensor, DistIndex]): def index_select(self, dist_index: Union[Tensor, DistIndex]):
......
starrygl/sample/batch_data.py
...@@ -7,6 +7,8 @@ from starrygl.sample.graph_core import DataSet ...@@ -7,6 +7,8 @@ from starrygl.sample.graph_core import DataSet
from starrygl.sample.graph_core import GraphData from starrygl.sample.graph_core import GraphData
from starrygl.sample.sample_core.base import BaseSampler, NegativeSampling from starrygl.sample.sample_core.base import BaseSampler, NegativeSampling
import dgl import dgl
from starrygl.sample.stream_manager import PipelineManager, getPipelineManger
""" """
入参不变，出参变为： 入参不变，出参变为：
sample_from_nodes sample_from_nodes
...@@ -42,7 +44,7 @@ def prepare_input(node_feat, edge_feat, mem_embedding,mfgs,dist_nid,dist_eid): ...@@ -42,7 +44,7 @@ def prepare_input(node_feat, edge_feat, mem_embedding,mfgs,dist_nid,dist_eid):
#print(idx.shape[0],b.srcdata['mem_ts'].shape) #print(idx.shape[0],b.srcdata['mem_ts'].shape)
return mfgs return mfgs
def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device = torch.device('cuda')): def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device = torch.device('cuda'),group = None):
if len(sample_out) > 1: if len(sample_out) > 1:
sample_out,metadata = sample_out sample_out,metadata = sample_out
...@@ -76,26 +78,25 @@ def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device = ...@@ -76,26 +78,25 @@ def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device =
dist_nid = nid_mapper[nid_tensor].to(device) dist_nid = nid_mapper[nid_tensor].to(device)
dist_nid,nid_inv = dist_nid.unique(return_inverse = True) dist_nid,nid_inv = dist_nid.unique(return_inverse = True)
if isinstance(graph.edge_attr,DistributedTensor): if isinstance(graph.edge_attr,DistributedTensor):
local_index, input_split,output_split = graph.edge_attr.gather_select_index(dist_eid) ind_dict = graph.edge_attr.all_to_all_ind2ptr(dist_eid,group = group)
edge_feat = graph.edge_attr.scatter_data(local_index,input_split=input_split,out_split=output_split) edge_feat = graph.edge_attr.all_to_all_get(group = group,**ind_dict)
else: else:
edge_feat = graph._get_edge_attr(dist_eid) edge_feat = graph._get_edge_attr(dist_eid)
local_index = None ind_dict = None
if isinstance(graph.x,DistributedTensor): if isinstance(graph.x,DistributedTensor):
local_index, input_split,output_split = graph.x.gather_select_index(dist_nid) ind_dict = graph.x.all_to_all_ind2ptr(dist_nid,group = group)
node_feat = graph.x.scatter_data(local_index,input_split=input_split,out_split=output_split) node_feat = graph.x.all_to_all_get(group = group,**ind_dict)
else: else:
node_feat = graph._get_node_attr(dist_nid) node_feat = graph._get_node_attr(dist_nid)
if mailbox is not None: if mailbox is not None:
if torch.distributed.get_world_size() > 1: if torch.distributed.get_world_size() > 1:
if node_feat is None: if node_feat is None:
local_index, input_split,output_split = mailbox.node_memory.gather_select_index(dist_nid) ind_dict = mailbox.node_memory.all_to_all_ind2ptr(dist_nid,group = group)
mem = mailbox.gather_memory(local_index,input_split,output_split) mem = mailbox.gather_memory(**ind_dict)
else: else:
mem = mailbox.get_memory(dist_nid) mem = mailbox.get_memory(dist_nid)
else: else:
mem = None mem = None
def build_block(): def build_block():
mfgs = list() mfgs = list()
col = torch.arange(0,root_len,device = device) col = torch.arange(0,root_len,device = device)
...@@ -110,23 +111,12 @@ def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device = ...@@ -110,23 +111,12 @@ def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device =
device = device) device = device)
idx = nid_inv[0:row_len + elen] idx = nid_inv[0:row_len + elen]
e_idx = eid_inv[col_len:col_len+elen] e_idx = eid_inv[col_len:col_len+elen]
b.srcdata['ID'] = idx#dist_nid[idx] b.srcdata['ID'] = idx
if sample_out[r].delta_ts().shape[0] > 0: if sample_out[r].delta_ts().shape[0] > 0:
b.edata['dt'] = sample_out[r].delta_ts().to(device) b.edata['dt'] = sample_out[r].delta_ts().to(device)
if src_ts is not None: if src_ts is not None:
b.srcdata['ts'] = src_ts[0:row_len + eid_len[r]] b.srcdata['ts'] = src_ts[0:row_len + eid_len[r]]
b.edata['ID'] = e_idx#dist_eid[e_idx] b.edata['ID'] = e_idx
#if edge_feat is not None:
# b.edata['f'] = edge_feat[e_idx]
#if r == len(eid_len)-1:
# if node_feat is not None:
# b.srcdata['h'] = node_feat[idx]
# if mem_embedding is not None:
# node_memory,node_memory_ts,mailbox,mailbox_ts = mem_embedding
# b.srcdata['mem'] = node_memory[idx]
# b.srcdata['mem_ts'] = node_memory_ts[idx]
# b.srcdata['mem_input'] = mailbox[idx].cuda().reshape(b.srcdata['ID'].shape[0], -1)
# b.srcdata['mail_ts'] = mailbox_ts[idx]
col = row col = row
col_len += eid_len[r] col_len += eid_len[r]
row_len += eid_len[r] row_len += eid_len[r]
...@@ -134,27 +124,7 @@ def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device = ...@@ -134,27 +124,7 @@ def to_block(graph: GraphData, data, sample_out, mailbox:MailBox = None,device =
mfgs = list(map(list, zip(*[iter(mfgs)]))) mfgs = list(map(list, zip(*[iter(mfgs)])))
mfgs.reverse() mfgs.reverse()
return data,mfgs,metadata return data,mfgs,metadata
data,mfgs,metadata = build_block() data,mfgs,metadata = build_block()
#if dist.get_world_size() > 1:
# if(node_feat is None):
# node_feat = torch.futures.Future()
# node_feat.set_result(None)
# if(edge_feat is None):
# edge_feat = torch.futures.Future()
# edge_feat.set_result(None)
# if(mem is None):
# mem = torch.futures.Future()
# mem.set_result(None)
# def callback(fs,mfgs,dist_nid,dist_eid):
# node_feat,edge_feat,mem_embedding = fs.value()
# node_feat = node_feat.value()
# edge_feat = edge_feat.value()
# mem_embedding = mem_embedding.value()
# return prepare_input(node_feat,edge_feat,mem_embedding,mfgs,dist_nid,dist_eid)
# cal = lambda fut: callback(fs=fut,mfgs = mfgs,dist_nid = dist_nid,dist_eid =dist_eid)
# return data,torch.futures.collect_all([node_feat,edge_feat,mem]).then(cal),metadata
#else:
mfgs = prepare_input(node_feat,edge_feat,mem,mfgs,dist_nid,dist_eid) mfgs = prepare_input(node_feat,edge_feat,mem,mfgs,dist_nid,dist_eid)
#return build_block(node_feat,edge_feat,mem)#data,mfgs,metadata #return build_block(node_feat,edge_feat,mem)#data,mfgs,metadata
return data,mfgs,metadata return data,mfgs,metadata
...@@ -164,9 +134,15 @@ def graph_sample(graph, sampler:BaseSampler, ...@@ -164,9 +134,15 @@ def graph_sample(graph, sampler:BaseSampler,
sample_fn, data, sample_fn, data,
neg_sampling = None, neg_sampling = None,
mailbox = None, mailbox = None,
device = torch.device('cuda')): device = torch.device('cuda'),
async_op = False):
out = sample_fn(sampler,data,neg_sampling) out = sample_fn(sampler,data,neg_sampling)
return to_block(graph,data,out,mailbox,device) if async_op == False:
return to_block(graph,data,out,mailbox,device)
else:
manger = getPipelineManger()
future = manger.submit('lookup',to_block,graph,data,out,mailbox,device)
return future
def sample_from_nodes(sampler:BaseSampler, data:DataSet, **kwargs): def sample_from_nodes(sampler:BaseSampler, data:DataSet, **kwargs):
out = sampler.sample_from_nodes(nodes=data.nodes.reshape(-1)) out = sampler.sample_from_nodes(nodes=data.nodes.reshape(-1))
......
starrygl/sample/cache/buffer_index.cpp 0 → 100644
#include <omp.h>
#include <torch/extension.h>
#include <time.h>
#include <random>
#include <phmap.h>
#include <boost/thread/mutex.hpp>
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#define MTX boost::mutex
using namespace std;
namespace py = pybind11;
namespace th = torch;
typedef int64_t NodeIDType;
typedef int64_t EdgeIDType;
typedef float WeightType;
typedef float TimeStampType;
#define EXTRAARGS , phmap::priv::hash_default_hash<K>, \
phmap::priv::hash_default_eq<K>, \
std::allocator<K>, 4, MTX
template <class K, class V>
using HashM = phmap::parallel_flat_hash_map<K, V EXTRAARGS>;
void init_cache_buffer(long long size){
}
at::Storage set_cache_buffer(at::Stroge ind){
}
//按策略更新替换cache上的数据，返回对应的索引
at::Storage update_buffer_index(string policy, at::Storge ind, u){
}
at::Storage get_buffer_index(at::Storage ind){
}
PYBIND11_MODULE(sample_cores, m)
{
m
.def("get_neighbors",
&get_neighbors,
py::return_value_policy::reference)
.
}
\ No newline at end of file
starrygl/sample/data_loader.py
...@@ -148,7 +148,7 @@ class DistributedDataLoader: ...@@ -148,7 +148,7 @@ class DistributedDataLoader:
return batch_data return batch_data
else : else :
raise StopIteration raise StopIteration
else : if self.queue_size > 0 :
num_reqs = min(self.queue_size - self.num_pending,self.expected_idx - self.submitted) num_reqs = min(self.queue_size - self.num_pending,self.expected_idx - self.submitted)
for _ in range(num_reqs): for _ in range(num_reqs):
if len(self.result_queue)>0: if len(self.result_queue)>0:
......
starrygl/sample/graph_core/__init__.py
...@@ -4,7 +4,7 @@ import os.path as osp ...@@ -4,7 +4,7 @@ import os.path as osp
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from torch_geometric.data import Data from torch_geometric.data import Data
class GraphData(): class GraphStore():
def __init__(self, pdata, device = torch.device('cuda'), all_on_gpu = False): def __init__(self, pdata, device = torch.device('cuda'), all_on_gpu = False):
self.device = device self.device = device
self.ids = pdata.ids.to(device) self.ids = pdata.ids.to(device)
...@@ -14,8 +14,8 @@ class GraphData(): ...@@ -14,8 +14,8 @@ class GraphData():
else: else:
self.edge_ts = None self.edge_ts = None
self.sample_graph = pdata.sample_graph self.sample_graph = pdata.sample_graph
self.nids_mapper = build_mapper(nids=pdata.ids.to(device)).data.to('cpu') self.nids_mapper = build_mapper(nids=pdata.ids.to(device)).dist.to('cpu')
self.eids_mapper = build_mapper(nids=pdata.eids.to(device)).data.to('cpu') self.eids_mapper = build_mapper(nids=pdata.eids.to(device)).dist.to('cpu')
if all_on_gpu: if all_on_gpu:
self.nids_mapper = self.nids_mapper.to(device) self.nids_mapper = self.nids_mapper.to(device)
self.eids_mapper = self.eids_mapper.to(device) self.eids_mapper = self.eids_mapper.to(device)
...@@ -30,9 +30,9 @@ class GraphData(): ...@@ -30,9 +30,9 @@ class GraphData():
self.x = None self.x = None
if hasattr(pdata,'edge_attr') and pdata.edge_attr is not None: if hasattr(pdata,'edge_attr') and pdata.edge_attr is not None:
if world_size > 1: if world_size > 1:
self.edge_attr = DistributedTensor(pdata.edge_attr.to('cpu').to(torch.float)) self.edge_attr = DistributedTensor(pdata.edge_attr.to('cuda').to(torch.float))
else: else:
self.edge_attr = pdata.edge_attr.to('cpu').to(torch.float) self.edge_attr = pdata.edge_attr.to('cuda').to(torch.float)
else: else:
self.edge_attr = None self.edge_attr = None
...@@ -106,7 +106,7 @@ class DataSet: ...@@ -106,7 +106,7 @@ class DataSet:
setattr(d,k,v[indx]) setattr(d,k,v[indx])
return d return d
class TemporalGraphData(): class TemporalGraphData(GraphStore):
def __init__(self,pdata,device): def __init__(self,pdata,device):
super(TemporalGraphData,self).__init__(pdata,device) super(TemporalGraphData,self).__init__(pdata,device)
def _set_temporal_batch_cache(self,size,pin_size): def _set_temporal_batch_cache(self,size,pin_size):
...@@ -117,7 +117,7 @@ class TemporalGraphData(): ...@@ -117,7 +117,7 @@ class TemporalGraphData():
class TemporalNeighborSampleGraph(GraphData): class TemporalNeighborSampleGraph(GraphStore):
def __init__(self, sample_graph=None, mode='full', eids_mapper=None): def __init__(self, sample_graph=None, mode='full', eids_mapper=None):
self.edge_index = sample_graph['edge_index'] self.edge_index = sample_graph['edge_index']
self.num_edges = self.edge_index.shape[1] self.num_edges = self.edge_index.shape[1]
......
starrygl/sample/memory/shared_mailbox.py
from typing import Union
from typing import List from typing import List
from typing import Optional
import torch import torch
from torch.distributed import rpc from torch.distributed import rpc
import torch_scatter import torch_scatter
...@@ -118,7 +120,7 @@ class SharedMailBox(): ...@@ -118,7 +120,7 @@ class SharedMailBox():
def set_mailbox_all_to_all(self,index,memory, def set_mailbox_all_to_all(self,index,memory,
memory_ts,mail,mail_ts, memory_ts,mail,mail_ts,
reduce_Op = None): reduce_Op = None,group = None):
#futs: List[torch.futures.Future] = [] #futs: List[torch.futures.Future] = []
if self.num_parts == 1: if self.num_parts == 1:
dist_index = DistIndex(index) dist_index = DistIndex(index)
...@@ -132,7 +134,7 @@ class SharedMailBox(): ...@@ -132,7 +134,7 @@ class SharedMailBox():
device = self.device) device = self.device)
indic = torch.searchsorted(index,self.partptr,right=False) indic = torch.searchsorted(index,self.partptr,right=False)
scatter_len_list = indic[1:] - indic[0:-1] scatter_len_list = indic[1:] - indic[0:-1]
torch.distributed.all_to_all_single(gather_len_list,scatter_len_list) torch.distributed.all_to_all_single(gather_len_list,scatter_len_list,group = group)
input_split = scatter_len_list.tolist() input_split = scatter_len_list.tolist()
output_split = gather_len_list.tolist() output_split = gather_len_list.tolist()
gather_id_list = torch.empty( gather_id_list = torch.empty(
...@@ -143,7 +145,7 @@ class SharedMailBox(): ...@@ -143,7 +145,7 @@ class SharedMailBox():
output_split = gather_len_list.tolist() output_split = gather_len_list.tolist()
torch.distributed.all_to_all_single( torch.distributed.all_to_all_single(
gather_id_list,index,output_split_sizes=output_split, gather_id_list,index,output_split_sizes=output_split,
input_split_sizes=input_split) input_split_sizes=input_split,group = group)
index = gather_id_list index = gather_id_list
gather_memory = torch.empty( gather_memory = torch.empty(
[gather_len_list.sum(),memory.shape[1]], [gather_len_list.sum(),memory.shape[1]],
...@@ -160,22 +162,81 @@ class SharedMailBox(): ...@@ -160,22 +162,81 @@ class SharedMailBox():
torch.distributed.all_to_all_single( torch.distributed.all_to_all_single(
gather_memory,memory, gather_memory,memory,
output_split_sizes=output_split, output_split_sizes=output_split,
input_split_sizes=input_split) input_split_sizes=input_split,group = group)
torch.distributed.all_to_all_single( torch.distributed.all_to_all_single(
gather_memory_ts,memory_ts, gather_memory_ts,memory_ts,
output_split_sizes=output_split, output_split_sizes=output_split,
input_split_sizes=input_split) input_split_sizes=input_split,group = group)
torch.distributed.all_to_all_single( torch.distributed.all_to_all_single(
gather_mail,mail, gather_mail,mail,
output_split_sizes=output_split, output_split_sizes=output_split,
input_split_sizes=input_split) input_split_sizes=input_split,group = group)
torch.distributed.all_to_all_single( torch.distributed.all_to_all_single(
gather_mail_ts,mail_ts, gather_mail_ts,mail_ts,
output_split_sizes=output_split, output_split_sizes=output_split,
input_split_sizes=input_split) input_split_sizes=input_split,group = group)
self.set_mailbox_local(DistIndex(index).loc,gather_mail,gather_mail_ts,Reduce_Op = reduce_Op) self.set_mailbox_local(DistIndex(index).loc,gather_mail,gather_mail_ts,Reduce_Op = reduce_Op)
self.set_memory_local(DistIndex(index).loc,gather_memory,gather_memory_ts, Reduce_Op = reduce_Op) self.set_memory_local(DistIndex(index).loc,gather_memory,gather_memory_ts, Reduce_Op = reduce_Op)
def set_mailbox_all_to_all(self,index,memory,
memory_ts,mail,mail_ts,
reduce_Op = None,group = None):
#futs: List[torch.futures.Future] = []
if self.num_parts == 1:
dist_index = DistIndex(index)
part_idx = dist_index.part
index = dist_index.loc
self.set_mailbox_local(index,mail,mail_ts)
self.set_memory_local(index,memory,memory_ts)
else:
gather_len_list = torch.empty([self.num_parts],
dtype = int,
device = self.device)
indic = torch.searchsorted(index,self.partptr,right=False)
scatter_len_list = indic[1:] - indic[0:-1]
torch.distributed.all_to_all_single(gather_len_list,scatter_len_list,group = group)
input_split = scatter_len_list.tolist()
output_split = gather_len_list.tolist()
gather_id_list = torch.empty(
[gather_len_list.sum()],
dtype = torch.long,
device = self.device)
input_split = scatter_len_list.tolist()
output_split = gather_len_list.tolist()
torch.distributed.all_to_all_single(
gather_id_list,index,output_split_sizes=output_split,
input_split_sizes=input_split,group = group)
index = gather_id_list
gather_memory = torch.empty(
[gather_len_list.sum(),memory.shape[1]],
dtype = memory.dtype,device = self.device)
gather_memory_ts = torch.empty(
[gather_len_list.sum()],
dtype = memory_ts.dtype,device = self.device)
gather_mail = torch.empty(
[gather_len_list.sum(),mail.shape[1]],
dtype = mail.dtype,device = self.device)
gather_mail_ts = torch.empty(
[gather_len_list.sum()],
dtype = mail_ts.dtype,device = self.device)
torch.distributed.all_to_all_single(
gather_memory,memory,
output_split_sizes=output_split,
input_split_sizes=input_split,group = group)
torch.distributed.all_to_all_single(
gather_memory_ts,memory_ts,
output_split_sizes=output_split,
input_split_sizes=input_split,group = group)
torch.distributed.all_to_all_single(
gather_mail,mail,
output_split_sizes=output_split,
input_split_sizes=input_split,group = group)
torch.distributed.all_to_all_single(
gather_mail_ts,mail_ts,
output_split_sizes=output_split,
input_split_sizes=input_split,group = group)
self.set_mailbox_local(DistIndex(index).loc,gather_mail,gather_mail_ts,Reduce_Op = reduce_Op)
self.set_memory_local(DistIndex(index).loc,gather_memory,gather_memory_ts, Reduce_Op = reduce_Op)
def get_update_mail(self,dist_indx_mapper, def get_update_mail(self,dist_indx_mapper,
src,dst,ts,edge_feats, src,dst,ts,edge_feats,
...@@ -234,9 +295,16 @@ class SharedMailBox(): ...@@ -234,9 +295,16 @@ class SharedMailBox():
#print(memory.shape[0]) #print(memory.shape[0])
return memory,memory_ts,mail,mail_ts return memory,memory_ts,mail,mail_ts
return torch.futures.collect_all([memory,memory_ts,mail,mail_ts]).then(callback) return torch.futures.collect_all([memory,memory_ts,mail,mail_ts]).then(callback)
def gather_memory(self,index,input_split,out_split): def gather_memory(
return self.node_memory.scatter_data(index,input_split,out_split),\ self,
self.node_memory_ts.scatter_data(index,input_split,out_split),\ dist_index: Union[torch.Tensor, DistIndex, None] = None,
self.mailbox.scatter_data(index,input_split,out_split),\ send_ptr: Optional[List[int]] = None,
self.mailbox_ts.scatter_data(index,input_split,out_split) recv_ptr: Optional[List[int]] = None,
recv_ind: Optional[List[int]] = None,
group = None
):
return self.node_memory.all_to_all_get(dist_index,send_ptr,recv_ptr,recv_ind,group),\
self.node_memory_ts.all_to_all_get(dist_index,send_ptr,recv_ptr,recv_ind,group),\
self.mailbox.all_to_all_get(dist_index,send_ptr,recv_ptr,recv_ind,group),\
self.mailbox_ts.all_to_all_get(dist_index,send_ptr,recv_ptr,recv_ind,group)
starrygl/sample/part_utils/partition_tgnn.py
...@@ -2,7 +2,8 @@ import parser ...@@ -2,7 +2,8 @@ import parser
from torch_sparse import SparseTensor from torch_sparse import SparseTensor
from torch_geometric.data import Data from torch_geometric.data import Data
from torch_geometric.utils import degree from torch_geometric.utils import degree
from .torch_utils import get_norm_temporal from starrygl.lib.libstarrygl_ops import metis_partition
from starrygl.lib.libstarrygl_ops_sampler import get_norm_temporal
import os.path as osp import os.path as osp
import os import os
import shutil import shutil
......
starrygl/sample/part_utils/setup.py deleted 100644 → 0
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
setup(
name='torch_utils',
ext_modules=[
CppExtension(
name='torch_utils',
sources=['./torch_utils.cpp'],
extra_compile_args=['-fopenmp','-Xlinker',' -export-dynamic'],
include_dirs=["../sample_core"],
),
],
cmdclass={
'build_ext': BuildExtension
})#
#setup(
# name='cpu_cache_manager',
# ext_modules=[
# CppExtension(
# name='cpu_cache_manager',
# sources=['cpu_cache_manager.cpp'],
# extra_compile_args=['-fopenmp','-Xlinker',' -export-dynamic'],
# include_dirs=["./"],
# ),
# ],
# cmdclass={
# 'build_ext': BuildExtension
# })#
#
\ No newline at end of file
starrygl/sample/stream_manager.py
from concurrent.futures import ThreadPoolExecutor, thread
from multiprocessing import Process
from multiprocessing.pool import ThreadPool
from typing import Deque
import torch import torch
class WorkStreamEvent: import torch.distributed as dist
def __init__(self): stage = ['train_stream','write_memory','write_mail','lookup']
self.train_stream = torch.cuda.Stream()
self.write_memory_stream = torch.cuda.Stream()
self.fetch_stream = torch.cuda.Stream() class PipelineManager:
self.write_mail_stream = torch.cuda.Stream() def __init__(self,num_tasks = 10):
self.event = None self.stream_set = {}
event = WorkStreamEvent() self.dist_set = {}
def get_event(): self.args_queue = {}
return event.event self.thread_pool = ThreadPoolExecutor(num_tasks)
for k in stage:
self.stream_set[k] = torch.cuda.Stream()
self.dist_set[k] = dist.new_group()
self.args_queue[k] = Deque()
def submit(self,state,func,*args):
future = self.thread_pool.submit(target=self.run, args=(state,func,args))
return future
def run(self,state,func,*args):
while torch.cuda.stream(self.stream_set[state]):
func(args, group = self.dist_set[state])
manger = None
def getPipelineManger():
global manger
if manger == None:
manger = PipelineManager()
return manger
\ No newline at end of file
test.sh 0 → 100644
nohup python train_tgnn.py --dataname TaoBao --world_size 4 --rank 0 > taobao4.out &
nohup python train_tgnn.py --dataname TaoBao --world_size 4 --rank 1 > taobao4.out &
nohup python train_tgnn.py --dataname TaoBao --world_size 4 --rank 2 > taobao4.out &
nohup python train_tgnn.py --dataname TaoBao --world_size 4 --rank 3 > taobao4.out &
nohup python train_tgnn.py --dataname GDELT --world_size 4 --rank 0 > GDELT4.out &
nohup python train_tgnn.py --dataname GDELT --world_size 4 --rank 1 > GDELT4.out &
nohup python train_tgnn.py --dataname GDELT --world_size 4 --rank 2 > GDELT4.out &
nohup python train_tgnn.py --dataname GDELT --world_size 4 --rank 3 > GDELT4.out &
nohup python train_tgnn.py --dataname ML25M --world_size 2 --rank 0 > ML25M2.out &
nohup python train_tgnn.py --dataname ML25M --world_size 2 --rank 1 > ML25M2.out &
nohup python train_tgnn.py --dataname TaoBao --world_size 2 --rank 0 > TaoBao2.out &
nohup python train_tgnn.py --dataname TaoBao --world_size 2 --rank 1 > TaoBao2.out &
nohup python train_tgnn.py --dataname ML25M --world_size 1 --rank 0 > ML25M1.out &
nohup python train_tgnn.py --dataname TaoBao --world_size 1 --rank 0 > TaoBao1.out &
nohup python train_tgnn.py --dataname GDELT --world_size 1 --rank 0 > GDELT1.out &
\ No newline at end of file
train_tgnn.py
...@@ -25,24 +25,7 @@ import os ...@@ -25,24 +25,7 @@ import os
from starrygl.sample.data_loader import DistributedDataLoader from starrygl.sample.data_loader import DistributedDataLoader
from starrygl.sample.batch_data import SAMPLE_TYPE from starrygl.sample.batch_data import SAMPLE_TYPE
"""
test command
python test.py --world_size 2 --rank 0
--world_size', default=4, type=int, metavar='W',
help='number of workers')
parser.add_argument('--rank', default=0, type=int, metavar='W',
help='rank of the worker')
parser.add_argument('--log_interval', type=int, default=10, metavar='N',
help='interval between training status logs')
parser.add_argument('--gamma', type=float, default=0.99, metavar='G',
help='how much to value future rewards')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed for reproducibility')
parser.add_argument('--num_sampler', type=int, default=10, metavar='S',
help='number of samplers')
parser.add_argument('--queue_size', type=int, default=10, metavar='S',
help='sampler queue size')
"""
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="RPC Reinforcement Learning Example", description="RPC Reinforcement Learning Example",
formatter_class=argparse.ArgumentDefaultsHelpFormatter, formatter_class=argparse.ArgumentDefaultsHelpFormatter,
...@@ -51,6 +34,8 @@ parser.add_argument('--rank', default=0, type=str, metavar='W', ...@@ -51,6 +34,8 @@ parser.add_argument('--rank', default=0, type=str, metavar='W',
help='name of dataset') help='name of dataset')
parser.add_argument('--world_size', default=1, type=int, metavar='W', parser.add_argument('--world_size', default=1, type=int, metavar='W',
help='number of negative samples') help='number of negative samples')
parser.add_argument('--dataname', default=1, type=str, metavar='W',
help='number of negative samples')
args = parser.parse_args() args = parser.parse_args()
from sklearn.metrics import average_precision_score, roc_auc_score from sklearn.metrics import average_precision_score, roc_auc_score
import torch import torch
...@@ -82,7 +67,7 @@ def main(): ...@@ -82,7 +67,7 @@ def main():
ctx = DistributedContext.init(backend="nccl", use_gpu=True) ctx = DistributedContext.init(backend="nccl", use_gpu=True)
device_id = torch.cuda.current_device() device_id = torch.cuda.current_device()
print('use cuda on',device_id) print('use cuda on',device_id)
pdata = partition_load("./dataset/here/GDELT", algo="metis_for_tgnn") pdata = partition_load("./dataset/here/{}".format(args.dataname), algo="metis_for_tgnn")
graph = GraphData(pdata = pdata) graph = GraphData(pdata = pdata)
sample_graph = TemporalNeighborSampleGraph(sample_graph = pdata.sample_graph,mode = 'full') sample_graph = TemporalNeighborSampleGraph(sample_graph = pdata.sample_graph,mode = 'full')
...@@ -174,10 +159,12 @@ def main(): ...@@ -174,10 +159,12 @@ def main():
dst = metadata['dst_pos_index'] dst = metadata['dst_pos_index']
ts = roots.ts ts = roots.ts
if(graph.edge_attr.device == torch.device('cpu')): if graph.edge_attr is None:
edge_feats = graph.edge_attr[roots.eids.to('cpu')].to('cuda') if graph.edge_attr is not None else None edge_feats = None
elif(graph.edge_attr.device == torch.device('cpu')):
edge_feats = graph.edge_attr[roots.eids.to('cpu')].to('cuda')
else: else:
edge_feats = graph.edge_attr[roots.eids] if graph.edge_attr is not None else None edge_feats = graph.edge_attr[roots.eids]
dist_index_mapper = mfgs[0][0].srcdata['ID'] dist_index_mapper = mfgs[0][0].srcdata['ID']
root_index = torch.cat((src,dst)) root_index = torch.cat((src,dst))
last_updated_nid = model.module.memory_updater.last_updated_nid[root_index] last_updated_nid = model.module.memory_updater.last_updated_nid[root_index]
...@@ -247,10 +234,12 @@ def main(): ...@@ -247,10 +234,12 @@ def main():
dst = metadata['dst_pos_index'] dst = metadata['dst_pos_index']
ts = roots.ts ts = roots.ts
if(graph.edge_attr.device == torch.device('cpu')): if graph.edge_attr is None:
edge_feats = graph.edge_attr[roots.eids.to('cpu')].to('cuda') if graph.edge_attr is not None else None edge_feats = None
elif(graph.edge_attr.device == torch.device('cpu')):
edge_feats = graph.edge_attr[roots.eids.to('cpu')].to('cuda')
else: else:
edge_feats = graph.edge_attr[roots.eids] if graph.edge_attr is not None else None edge_feats = graph.edge_attr[roots.eids]
dist_index_mapper = mfgs[0][0].srcdata['ID'] dist_index_mapper = mfgs[0][0].srcdata['ID']
root_index = torch.cat((src,dst)) root_index = torch.cat((src,dst))
last_updated_nid = model.module.memory_updater.last_updated_nid[root_index] last_updated_nid = model.module.memory_updater.last_updated_nid[root_index]
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment