Skip to content

Commit 2eb25c4

Browse files
SolenoidWGTSolenoidWGT
committed
feature(wgt): enable DI using torch-rpc to support GPU-p2p and RDMA-rpc
1. Add torchrpc message queue. 2. Implement buffer based on CUDA-shared-tensor to optimize the data path of torchrpc. 3. Add 'bypass_eventloop' arg in Task() and Parallel(). 4. Add thread lock in distributer.py to prevent sender and receiver competition. 5. Add message queue perf test for torchrpc, nccl, nng, shm 6. Add comm_perf_helper.py to make program timing more convenient. 7. Modified the subscribe() of class MQ, adding 'fn' parameter and 'is_once' parameter. 8. Add new DummyLock and ConditionLock type in lock_helper.py 9. Add message queues perf test. 10. Introduced a new self-hosted runner to execute cuda, multiprocess, torchrpc related tests.
1 parent c30818a commit 2eb25c4

34 files changed

+2356
-123
lines changed

‎.github/workflows/unit_test.yml‎

Lines changed: 75 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,11 @@ jobs:
1111
if: "!contains(github.event.head_commit.message, 'ci skip')"
1212
strategy:
1313
matrix:
14-
python-version: [3.7, 3.8, 3.9]
15-
14+
python-version: ["3.7", "3.8", "3.9"]
1615
steps:
17-
- uses: actions/checkout@v2
16+
- uses: actions/checkout@v3
1817
- name: Set up Python ${{ matrix.python-version }}
19-
uses: actions/setup-python@v2
18+
uses: actions/setup-python@v3
2019
with:
2120
python-version: ${{ matrix.python-version }}
2221
- name: do_unittest
@@ -41,12 +40,13 @@ jobs:
4140
if: "!contains(github.event.head_commit.message, 'ci skip')"
4241
strategy:
4342
matrix:
44-
python-version: [3.7, 3.8, 3.9]
45-
43+
python-version: ["3.7", "3.8", "3.9"]
4644
steps:
47-
- uses: actions/checkout@v2
45+
- uses: actions/checkout@v3
4846
- name: Set up Python ${{ matrix.python-version }}
49-
uses: actions/setup-python@v2
47+
uses: actions/setup-python@v3
48+
env:
49+
AGENT_TOOLSDIRECTORY: /opt/hostedtoolcache
5050
with:
5151
python-version: ${{ matrix.python-version }}
5252
- name: do_benchmark
@@ -55,3 +55,70 @@ jobs:
5555
python -m pip install ".[test,k8s]"
5656
./ding/scripts/install-k8s-tools.sh
5757
make benchmark
58+
59+
test_multiprocess:
60+
runs-on: self-hosted
61+
if: "!contains(github.event.head_commit.message, 'ci skip')"
62+
strategy:
63+
matrix:
64+
python-version: ["3.7", "3.8", "3.9"]
65+
steps:
66+
- uses: actions/checkout@v3
67+
- name: Set up Python ${{ matrix.python-version }}
68+
uses: actions/setup-python@v3
69+
with:
70+
python-version: ${{ matrix.python-version }}
71+
- name: do_multiprocesstest
72+
timeout-minutes: 40
73+
run: |
74+
python -m pip install box2d-py
75+
python -m pip install .
76+
python -m pip install ".[test,k8s]"
77+
./ding/scripts/install-k8s-tools.sh
78+
make multiprocesstest
79+
80+
test_cuda:
81+
runs-on: self-hosted
82+
if: "!contains(github.event.head_commit.message, 'ci skip')"
83+
strategy:
84+
matrix:
85+
python-version: ["3.7", "3.8", "3.9"]
86+
steps:
87+
- uses: actions/checkout@v3
88+
- name: Set up Python ${{ matrix.python-version }}
89+
uses: actions/setup-python@v3
90+
env:
91+
AGENT_TOOLSDIRECTORY: /opt/hostedtoolcache
92+
with:
93+
python-version: ${{ matrix.python-version }}
94+
- name: do_unittest
95+
timeout-minutes: 40
96+
run: |
97+
python -m pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
98+
python -m pip install box2d-py
99+
python -m pip install .
100+
python -m pip install ".[test,k8s]"
101+
./ding/scripts/install-k8s-tools.sh
102+
make cudatest
103+
104+
test_mq_benchmark:
105+
runs-on: self-hosted
106+
if: "!contains(github.event.head_commit.message, 'ci skip')"
107+
strategy:
108+
matrix:
109+
python-version: ["3.7", "3.8", "3.9"]
110+
steps:
111+
- uses: actions/checkout@v3
112+
- name: Set up Python ${{ matrix.python-version }}
113+
uses: actions/setup-python@v3
114+
env:
115+
AGENT_TOOLSDIRECTORY: /opt/hostedtoolcache
116+
with:
117+
python-version: ${{ matrix.python-version }}
118+
- name: do_mqbenchmark
119+
run: |
120+
python -m pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
121+
python -m pip install .
122+
python -m pip install ".[test,k8s]"
123+
./ding/scripts/install-k8s-tools.sh
124+
make mqbenchmark

‎Makefile‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,11 +57,25 @@ benchmark:
5757
--durations=0 \
5858
-sv -m benchmark
5959

60+
multiprocesstest:
61+
pytest ${TEST_DIR} \
62+
--cov-report=xml \
63+
--cov-report term-missing \
64+
--cov=${COV_DIR} \
65+
${DURATIONS_COMMAND} \
66+
${WORKERS_COMMAND} \
67+
-sv -m multiprocesstest
68+
69+
mqbenchmark:
70+
pytest ${TEST_DIR} \
71+
--durations=0 \
72+
-sv -m mqbenchmark
73+
6074
test: unittest # just for compatibility, can be changed later
6175

6276
cpu_test: unittest algotest benchmark
6377

64-
all_test: unittest algotest cudatest benchmark
78+
all_test: unittest algotest cudatest benchmark multiprocesstest
6579

6680
format:
6781
yapf --in-place --recursive -p --verbose --style .style.yapf ${FORMAT_DIR}

‎codecov.yml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,10 @@ coverage:
66
target: auto
77
threshold: 0.5%
88
if_ci_failed: success #success, failure, error, ignore
9+
10+
# fix me
11+
# The unittests of the torchrpc module are tested by different runners and cannot be included
12+
# in the test_unittest's coverage report. To keep CI happy, we don't count torchrpc related coverage.
13+
ignore:
14+
- /mnt/cache/wangguoteng/DI-engine/ding/framework/message_queue/torch_rpc.py
15+
- /mnt/cache/wangguoteng/DI-engine/ding/framework/message_queue/perfs/*

‎ding/compatibility.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,7 @@ def torch_ge_131():
77

88
def torch_ge_180():
99
return int("".join(list(filter(str.isdigit, torch.__version__)))) >= 180
10+
11+
12+
def torch_ge_1121():
13+
return int("".join(list(filter(str.isdigit, torch.__version__)))) >= 1121

‎ding/data/shm_buffer.py‎

Lines changed: 130 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,10 @@
33
import ctypes
44
import numpy as np
55
import torch
6+
import torch.multiprocessing as mp
7+
from functools import reduce
8+
from ditk import logging
9+
from abc import abstractmethod
610

711
_NTYPE_TO_CTYPE = {
812
np.bool_: ctypes.c_bool,
@@ -18,8 +22,37 @@
1822
np.float64: ctypes.c_double,
1923
}
2024

25+
# uint16, uint32, uint32
26+
_NTYPE_TO_TTYPE = {
27+
np.bool_: torch.bool,
28+
np.uint8: torch.uint8,
29+
# np.uint16: torch.int16,
30+
# np.uint32: torch.int32,
31+
# np.uint64: torch.int64,
32+
np.int8: torch.uint8,
33+
np.int16: torch.int16,
34+
np.int32: torch.int32,
35+
np.int64: torch.int64,
36+
np.float32: torch.float32,
37+
np.float64: torch.float64,
38+
}
39+
40+
_NOT_SUPPORT_NTYPE = {np.uint16: torch.int16, np.uint32: torch.int32, np.uint64: torch.int64}
41+
_CONVERSION_TYPE = {np.uint16: np.int16, np.uint32: np.int32, np.uint64: np.int64}
42+
43+
44+
class ShmBufferBase:
45+
46+
@abstractmethod
47+
def fill(self, src_arr: Union[np.ndarray, torch.Tensor]) -> None:
48+
raise NotImplementedError
2149

22-
class ShmBuffer():
50+
@abstractmethod
51+
def get(self) -> Union[np.ndarray, torch.Tensor]:
52+
raise NotImplementedError
53+
54+
55+
class ShmBuffer(ShmBufferBase):
2356
"""
2457
Overview:
2558
Shared memory buffer to store numpy array.
@@ -78,6 +111,94 @@ def get(self) -> np.ndarray:
78111
return data
79112

80113

114+
class ShmBufferCuda(ShmBufferBase):
115+
116+
def __init__(
117+
self,
118+
dtype: Union[torch.dtype, np.dtype],
119+
shape: Tuple[int],
120+
ctype: Optional[type] = None,
121+
copy_on_get: bool = True,
122+
device: Optional[torch.device] = torch.device('cuda:0')
123+
) -> None:
124+
"""
125+
Overview:
126+
Use torch.multiprocessing for shared tensor or ndaray between processes.
127+
Arguments:
128+
- dtype (Union[torch.dtype, np.dtype]): dtype of torch.tensor or numpy.ndarray.
129+
- shape (Tuple[int]): Shape of torch.tensor or numpy.ndarray.
130+
- ctype (type): Origin class type, e.g. np.ndarray, torch.Tensor.
131+
- copy_on_get (bool, optional): Can be set to False only if the shared object
132+
is a tenor, otherwise True.
133+
- device (Optional[torch.device], optional): The GPU device where cuda-shared-tensor
134+
is located, the default is cuda:0.
135+
136+
Raises:
137+
RuntimeError: Unsupported share type by ShmBufferCuda.
138+
"""
139+
if isinstance(dtype, np.dtype): # it is type of gym.spaces.dtype
140+
self.ctype = np.ndarray
141+
dtype = dtype.type
142+
if dtype in _NOT_SUPPORT_NTYPE.keys():
143+
logging.warning(
144+
"Torch tensor unsupport numpy type {}, attempt to do a type conversion, which may lose precision.".
145+
format(dtype)
146+
)
147+
ttype = _NOT_SUPPORT_NTYPE[dtype]
148+
self.dtype = _CONVERSION_TYPE[dtype]
149+
else:
150+
ttype = _NTYPE_TO_TTYPE[dtype]
151+
self.dtype = dtype
152+
elif isinstance(dtype, torch.dtype):
153+
self.ctype = torch.Tensor
154+
ttype = dtype
155+
else:
156+
raise RuntimeError("The dtype parameter only supports torch.dtype and np.dtype")
157+
158+
self.copy_on_get = copy_on_get
159+
self.shape = shape
160+
self.device = device
161+
# We don't want the buffer to be involved in the computational graph
162+
with torch.no_grad():
163+
self.buffer = torch.zeros(reduce(lambda x, y: x * y, shape), dtype=ttype, device=self.device)
164+
165+
def fill(self, src_arr: Union[np.ndarray, torch.Tensor]) -> None:
166+
if self.ctype is np.ndarray:
167+
if src_arr.dtype.type != self.dtype:
168+
logging.warning(
169+
"Torch tensor unsupport numpy type {}, attempt to do a type conversion, which may lose precision.".
170+
format(self.dtype)
171+
)
172+
src_arr = src_arr.astype(self.dtype)
173+
tensor = torch.from_numpy(src_arr)
174+
elif self.ctype is torch.Tensor:
175+
tensor = src_arr
176+
else:
177+
raise RuntimeError("Unsopport CUDA-shared-tensor input type:\"{}\"".format(type(src_arr)))
178+
179+
# If the GPU-a and GPU-b are connected using nvlink, the copy is very fast.
180+
with torch.no_grad():
181+
self.buffer.copy_(tensor.view(tensor.numel()))
182+
183+
def get(self) -> Union[np.ndarray, torch.Tensor]:
184+
with torch.no_grad():
185+
if self.ctype is np.ndarray:
186+
# Because ShmBufferCuda use CUDA memory exchanging data between processes.
187+
# So copy_on_get is necessary for numpy arrays.
188+
re = self.buffer.cpu()
189+
re = re.detach().view(self.shape).numpy()
190+
else:
191+
if self.copy_on_get:
192+
re = self.buffer.clone().detach().view(self.shape)
193+
else:
194+
re = self.buffer.view(self.shape)
195+
196+
return re
197+
198+
def __del__(self):
199+
del self.buffer
200+
201+
81202
class ShmBufferContainer(object):
82203
"""
83204
Overview:
@@ -88,7 +209,8 @@ def __init__(
88209
self,
89210
dtype: Union[Dict[Any, type], type, np.dtype],
90211
shape: Union[Dict[Any, tuple], tuple],
91-
copy_on_get: bool = True
212+
copy_on_get: bool = True,
213+
is_cuda_buffer: bool = False
92214
) -> None:
93215
"""
94216
Overview:
@@ -98,11 +220,15 @@ def __init__(
98220
- shape (:obj:`Union[Dict[Any, tuple], tuple]`): If `Dict[Any, tuple]`, use a dict to manage \
99221
multiple buffers; If `tuple`, use single buffer.
100222
- copy_on_get (:obj:`bool`): Whether to copy data when calling get method.
223+
- is_cuda_buffer (:obj:`bool`): Whether to use pytorch CUDA shared tensor as the implementation of shm.
101224
"""
102225
if isinstance(shape, dict):
103-
self._data = {k: ShmBufferContainer(dtype[k], v, copy_on_get) for k, v in shape.items()}
226+
self._data = {k: ShmBufferContainer(dtype[k], v, copy_on_get, is_cuda_buffer) for k, v in shape.items()}
104227
elif isinstance(shape, (tuple, list)):
105-
self._data = ShmBuffer(dtype, shape, copy_on_get)
228+
if not is_cuda_buffer:
229+
self._data = ShmBuffer(dtype, shape, copy_on_get)
230+
else:
231+
self._data = ShmBufferCuda(dtype, shape, copy_on_get)
106232
else:
107233
raise RuntimeError("not support shape: {}".format(shape))
108234
self._shape = shape

0 commit comments

Comments
 (0)