├── LICENSE ├── README.md ├── rdma_test.py └── utils ├── __init__.py ├── __pycache__ ├── __init__.cpython-38.pyc ├── connection.cpython-38.pyc └── param_parser.cpython-38.pyc ├── connection.py └── param_parser.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Li Weihang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # python_rdma_test 2 | 3 | This is an RDMA program written in Python, based on the Pyverbs provided by the rdma-core(https://github.com/linux-rdma/rdma-core) repository. 4 | 5 | 6 | 7 | # How to run 8 | 9 | 1. Clone the rdma-core repository. 10 | 11 | 2. Follow README.md of rdma-core and then build the project, please make sure pyverbs is compliled successfully. 12 | 13 | 3. Set PYTHONPATH to let the Python interpreter find where Pyverbs is. 14 | 15 | 4. Run rdma_test.py. 16 | 17 | 18 | 19 | # Example 20 | 21 | 1. Show help 22 | ```bash 23 | PYTHONPATH=../rdma-core/build/python/ ./rdma.py -h 24 | ``` 25 | 26 | 2. Run RDMA Write between two nodes with RC QP: 27 | - Server 28 | ```bash 29 | PYTHONPATH=../rdma-core/build/python/ ./rdma.py -d rxe_0 -o write 30 | ``` 31 | 32 | - Client 33 | ```bash 34 | PYTHONPATH=../rdma-core/build/python/ ./rdma.py -d rxe_0 -o write 192.168.xx.xx 35 | ``` 36 | -------------------------------------------------------------------------------- /rdma_test.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3.8 2 | import sys 3 | 4 | sys.path.append('..') 5 | 6 | from utils.connection import SKT, CM 7 | from utils.param_parser import parser 8 | 9 | from pyverbs.addr import AH, AHAttr, GlobalRoute 10 | from pyverbs.cq import CQ 11 | from pyverbs.device import Context 12 | from pyverbs.enums import * 13 | from pyverbs.mr import MR 14 | from pyverbs.pd import PD 15 | from pyverbs.qp import QP, QPCap, QPInitAttr, QPAttr 16 | from pyverbs.wr import SGE, RecvWR, SendWR 17 | 18 | 19 | RECV_WR = 1 20 | SEND_WR = 2 21 | GRH_LENGTH = 40 22 | 23 | # TODO: Error handling 24 | 25 | args = parser.parse_args() 26 | 27 | server = not bool(args['server_ip']) 28 | 29 | if args['use_cm']: 30 | conn = CM(args['port'], args['server_ip']) 31 | else: 32 | conn = SKT(args['port'], args['server_ip']) 33 | 34 | print('-' * 80) 35 | print(' ' * 25, "Python test for RDMA") 36 | 37 | if server: 38 | print("Running as server...") 39 | else: 40 | print("Running as client...") 41 | 42 | print('-' * 80) 43 | 44 | if args['qp_type'] == IBV_QPT_UD and args['operation_type'] != IBV_WR_SEND: 45 | print("UD QPs don't support RDMA operations.") 46 | conn.close() 47 | 48 | conn.handshake() 49 | 50 | ctx = Context(name=args['ib_dev']) 51 | pd = PD(ctx) 52 | cq = CQ(ctx, 100) 53 | 54 | cap = QPCap(max_send_wr=args['tx_depth'], max_recv_wr=args['rx_depth'], max_send_sge=args['sg_depth'], 55 | max_recv_sge=args['sg_depth'], max_inline_data=args['inline_size']) 56 | qp_init_attr = QPInitAttr(qp_type=args['qp_type'], scq=cq, rcq=cq, cap=cap, sq_sig_all=True) 57 | qp = QP(pd, qp_init_attr) 58 | 59 | gid = ctx.query_gid(port_num=1, index=args['gid_index']) 60 | 61 | # Handshake to exchange information such as QP Number 62 | remote_info = conn.handshake(gid=gid, qpn=qp.qp_num) 63 | 64 | gr = GlobalRoute(dgid=remote_info['gid'], sgid_index=args['gid_index']) 65 | ah_attr = AHAttr(gr=gr, is_global=1, port_num=1) 66 | 67 | if args['qp_type'] == IBV_QPT_UD: 68 | ah = AH(pd, attr=ah_attr) 69 | qp.to_rts(QPAttr()) 70 | else: 71 | qa = QPAttr() 72 | qa.ah_attr = ah_attr 73 | qa.dest_qp_num = remote_info['qpn'] 74 | qa.path_mtu = args['mtu'] 75 | qa.max_rd_atomic = 1 76 | qa.max_dest_rd_atomic = 1 77 | qa.qp_access_flags = IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_LOCAL_WRITE 78 | if server: 79 | qp.to_rtr(qa) 80 | else: 81 | qp.to_rts(qa) 82 | 83 | conn.handshake() 84 | 85 | mr_size = args['size'] 86 | if server: 87 | if args['qp_type'] == IBV_QPT_UD: # UD needs more space to store GRH when receiving. 88 | mr_size = mr_size + GRH_LENGTH 89 | content = 's' * mr_size 90 | else: 91 | content = 'c' * mr_size 92 | 93 | mr = MR(pd, mr_size, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ) 94 | sgl = [SGE(mr.buf, mr.length, mr.lkey)] 95 | 96 | if args['operation_type'] != IBV_WR_SEND: 97 | remote_info = conn.handshake(addr=mr.buf, rkey=mr.rkey) 98 | 99 | def read_mr(mr): 100 | if args['qp_type'] == IBV_QPT_UD and server: 101 | return mr.read(mr.length - GRH_LENGTH, GRH_LENGTH).decode() 102 | else: 103 | return mr.read(mr.length, 0).decode() 104 | 105 | for i in range(args['iters']): 106 | print("Iter: " + f"{i + 1}/{args['iters']}") 107 | 108 | mr.write(content, len(content)) 109 | print("MR Content before test:" + read_mr(mr)) 110 | 111 | if server and args['operation_type'] == IBV_WR_SEND: 112 | wr = RecvWR(RECV_WR, len(sgl), sgl) 113 | qp.post_recv(wr) 114 | 115 | conn.handshake() 116 | 117 | if not server: 118 | wr = SendWR(SEND_WR, opcode=args['operation_type'], num_sge=1, sg=sgl) 119 | if args['qp_type'] == IBV_QPT_UD: 120 | wr.set_wr_ud(ah, remote_info['qpn'], 0) 121 | elif args['operation_type'] != IBV_WR_SEND: 122 | wr.set_wr_rdma(remote_info['rkey'], remote_info['addr']) 123 | 124 | qp.post_send(wr) 125 | 126 | conn.handshake() 127 | 128 | if not server or args['operation_type'] == IBV_WR_SEND: 129 | wc_num, wc_list = cq.poll() 130 | 131 | print("MR Content after test:" + read_mr(mr)) 132 | 133 | conn.handshake() 134 | conn.close() 135 | 136 | print('-' * 80) -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.insert(0, '../../build/python') -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Li-Weihang/python_rdma_test/cc3bd8b17508724d77c418823bded49170cd2008/utils/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/connection.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Li-Weihang/python_rdma_test/cc3bd8b17508724d77c418823bded49170cd2008/utils/__pycache__/connection.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/param_parser.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Li-Weihang/python_rdma_test/cc3bd8b17508724d77c418823bded49170cd2008/utils/__pycache__/param_parser.cpython-38.pyc -------------------------------------------------------------------------------- /utils/connection.py: -------------------------------------------------------------------------------- 1 | import socket 2 | import time 3 | # All supported types should be imported here for the parser of received message 4 | # or the Global() method can't get objects' type from the received string 5 | from builtins import str, int 6 | from pyverbs.addr import GID 7 | 8 | from pyverbs.cm_enums import * 9 | from pyverbs.cmid import CMID, AddrInfo 10 | from pyverbs.qp import QPInitAttr, QPCap 11 | 12 | RESERVED_LEN = 20 # We should reserve some space for receiving because the length of exchanged strings may be not equal 13 | HANDSHAKE_WORDS = "Aloha!" 14 | 15 | 16 | class CommError(Exception): 17 | def __init__(self, message): 18 | self.message = message 19 | super(CommError, self).__init__(message) 20 | 21 | def __str__(self): 22 | return self.message 23 | 24 | 25 | class CommBase: 26 | def __init__(self): 27 | pass 28 | 29 | @staticmethod 30 | def prepare_send_msg(**kwargs): 31 | if not bool(kwargs): 32 | send_msg = HANDSHAKE_WORDS 33 | else: 34 | print("-- Local Info") 35 | # Prepare the info to be sent 36 | send_msg = "" 37 | for key, value in kwargs.items(): 38 | print(key, ":", value) 39 | send_msg = send_msg + key + ':' + type(value).__name__ + ':' + str(value) + ',' 40 | 41 | send_msg = send_msg[:-1] 42 | print('-' * 80) 43 | 44 | return send_msg 45 | 46 | @staticmethod 47 | def parse_recv_msg(only_handshake, recv_msg): 48 | if only_handshake: 49 | try: 50 | if recv_msg != HANDSHAKE_WORDS: 51 | raise CommError("Failed to handshake with remote peer by " + CommBase.__class__.__name__) 52 | except CommError as e: 53 | print(e) 54 | else: 55 | # Parse the recv info, create objects in specified type and organize them into a dictionary 56 | key_value = {} 57 | print("-- Remote Info") 58 | for item in recv_msg.split(','): 59 | key, value_type, value = item.split(':', 2) 60 | key_value[key.strip()] = globals()[value_type.strip()](value.strip()) 61 | print(key, ":", value) 62 | print('-' * 80) 63 | 64 | return key_value 65 | 66 | 67 | # TODO: Error Handling/Handshake Overtime 68 | class SKT(CommBase): 69 | def __init__(self, port, ip=None): 70 | """ 71 | Initializes a Socket object for RDMA test, the process of socket connection has 72 | been encapsulated in this constructor 73 | :param port: The port number for TCP/IP socket 74 | :param ip: The ip address of the server, only the client needs this param 75 | """ 76 | super(CommBase, self).__init__() 77 | self.socket = socket.socket() 78 | # Avoid the "Address already in use" error 79 | self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 80 | self.isServer = False 81 | 82 | if ip is None: # is Server 83 | self.isServer = True 84 | self.socket.bind(('0.0.0.0', port)) # All IP can be accepted. 85 | self.socket.listen() 86 | self.client, self.addr = self.socket.accept() 87 | self.socket.close() # We only support single client, no more clients will be accepted 88 | self.socket = self.client 89 | else: 90 | self.socket.connect((ip, port)) 91 | 92 | def handshake(self, **kwargs): 93 | """ 94 | Makes a handshake between the server and client. The users can exchange anything the want with the remote peer. 95 | For example: 96 | 1. socket.handshake() 97 | This exchanges nothing, which is always used as a synchronization mechanism. 98 | 2. socket.handshake(qpn = xx, psn = yy) 99 | This exchanges qpn and psn with remote peer, this can be used before INIT->RTR. 100 | 3. socket.handshake(str = "hello bro") 101 | This exchanges a message between server and client. 102 | :param kwargs: Users should pass args in the form of 'key1 = value1, key2 = value2 ...', and both ends should 103 | provide same keys. 104 | :return: A dictionary based on received info if any param is passed in, or nothing. 105 | """ 106 | just_handshake = not bool(kwargs) 107 | 108 | send_msg = self.prepare_send_msg(**kwargs) 109 | 110 | self.socket.send(send_msg.encode('utf-8')) 111 | 112 | recv_msg = self.socket.recv(len(send_msg) + RESERVED_LEN).decode('utf-8') 113 | 114 | time.sleep(0.5) # wait for a while in case that the program can't spilt message between handshakes. 115 | 116 | return self.parse_recv_msg(just_handshake, recv_msg) 117 | 118 | def close(self): 119 | if not self.isServer: 120 | self.socket.close() 121 | 122 | def __del__(self): 123 | self.socket.close() 124 | 125 | 126 | # TODO: Error Handling 127 | class CM(CommBase): 128 | def __init__(self, port, ip=None): 129 | """ 130 | Initializes a CMID object for RDMA test, the process of creating CM connection has 131 | been encapsulated in this constructor 132 | :param port: The port number for TCP/IP socket 133 | :param ip: The ip address of the server, only the client needs this param 134 | """ 135 | super(CommBase, self).__init__() 136 | cap = QPCap(max_recv_wr=1) 137 | qp_init_attr = QPInitAttr(cap=cap) 138 | self.isServer = False 139 | 140 | # TODO: Create a QP after connection instead of creating one each time user want to handshake 141 | 142 | if ip is None: 143 | self.isServer = True 144 | cai = AddrInfo(src='0.0.0.0', src_service=str(port), port_space=RDMA_PS_TCP, flags=RAI_PASSIVE) 145 | self.cmid = CMID(creator=cai, qp_init_attr=qp_init_attr) 146 | self.cmid.listen() 147 | client_cmid = self.cmid.get_request() 148 | client_cmid.accept() 149 | self.cmid.close() # # We only support single client, no more clients will be accepted 150 | self.cmid = client_cmid 151 | else: 152 | cai = AddrInfo(src='0.0.0.0', dst=ip, dst_service=str(port), port_space=RDMA_PS_TCP, flags=0) 153 | self.cmid = CMID(creator=cai, qp_init_attr=qp_init_attr) 154 | self.cmid.connect() 155 | 156 | def handshake(self, **kwargs): 157 | """ 158 | Makes a handshake between the server and client. The users can exchange anything the want with the remote peer. 159 | For example: 160 | 1. socket.handshake() 161 | This exchanges nothing, which is always used as a synchronization mechanism. 162 | 2. socket.handshake(qpn = xx, psn = yy) 163 | This exchanges qpn and psn with remote peer, this can be used before INIT->RTR. 164 | 3. socket.handshake(str = "hello bro") 165 | This exchanges a message between server and client. 166 | :param kwargs: Users should pass args in the form of 'key1 = value1, key2 = value2 ...', and both ends should 167 | provide same keys. 168 | :return: A dictionary based on received info if any param is passed in, or nothing. 169 | """ 170 | just_handshake = not bool(kwargs) 171 | 172 | send_msg = self.prepare_send_msg(**kwargs) 173 | size = len(send_msg) 174 | 175 | # Recv Message 176 | recv_mr = self.cmid.reg_msgs(size + RESERVED_LEN) 177 | self.cmid.post_recv(recv_mr, size) 178 | 179 | # Send Message 180 | send_mr = self.cmid.reg_msgs(size) 181 | send_mr.write(send_msg.encode('utf-8'), size) 182 | self.cmid.post_send(send_mr, 0, size) 183 | 184 | recv_wc = self.cmid.get_recv_comp() 185 | send_wc = self.cmid.get_send_comp() 186 | 187 | recv_msg = recv_mr.read(recv_wc.byte_len, 0).decode('utf-8') 188 | # print(recv_msg) 189 | 190 | return self.parse_recv_msg(just_handshake, recv_msg) 191 | 192 | def close(self): 193 | if not self.isServer: 194 | self.cmid.close() 195 | 196 | def __del__(self): 197 | self.cmid.close() -------------------------------------------------------------------------------- /utils/param_parser.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | 4 | from pyverbs.enums import IBV_QPT_UD, IBV_QPT_RC, IBV_WR_SEND, IBV_WR_RDMA_WRITE, IBV_WR_RDMA_READ 5 | 6 | 7 | class DictAction(argparse.Action): 8 | def __call__(self, parsers, namespace, values, option_string=None): 9 | setattr(namespace, self.dest, self.choices.get(values)) 10 | 11 | 12 | qp_dict = {'rc': IBV_QPT_RC, 'ud': IBV_QPT_UD} 13 | op_dict = {'send': IBV_WR_SEND, 'write': IBV_WR_RDMA_WRITE, 'read': IBV_WR_RDMA_READ} 14 | 15 | 16 | class ArgsParser(object): 17 | def __init__(self): 18 | self.args = None 19 | 20 | @staticmethod 21 | def parse_args(): 22 | arg_parser = argparse.ArgumentParser(description="A Python Test Program for RDMA.") 23 | arg_parser.add_argument('server_ip', type=str, nargs='?', 24 | help='The IP address of the Server (Client only).') 25 | arg_parser.add_argument('-C', '--use_cm', action='store_true', 26 | help='Use Connection Management protocol to handshake.') 27 | arg_parser.add_argument('-d', '--ib_dev', required=True, 28 | help='RDMA device to run the tests on.') 29 | arg_parser.add_argument('-G', '--sg_depth', type=int, default=1, 30 | help='Number of sge in the sgl.') 31 | arg_parser.add_argument('-I', '--inline_size', type=int, default=0, 32 | help='Max size of message to be sent in inline.') 33 | arg_parser.add_argument('-m', '--mtu', type=int, default=4, choices=range(0, 5), 34 | help='The Path MTU 0 ~ 4 (256/512/1024/2048/4096).') 35 | arg_parser.add_argument('-n', '--iters', type=int, default=1, 36 | help='Number of exchanges (at least 5, default 1).') 37 | arg_parser.add_argument('-p', '--port', type=int, default=18515, 38 | help='Listen on/connect to port (default 18515).') 39 | arg_parser.add_argument('-r', '--rx_depth', type=int, default=16, 40 | help='Size of rx queue (default 16). If using srq, rx-depth controls max-wr size of the srq.') 41 | arg_parser.add_argument('-s', '--size', type=int, default=13, 42 | help='Size of message to exchange (default 13).') 43 | arg_parser.add_argument('-S', '--sl', type=int, default=0, 44 | help='Service Level (default 0).') 45 | arg_parser.add_argument('-o', '--operation_type', default=IBV_WR_SEND, choices=op_dict, action=DictAction, 46 | help='The type of operation.') 47 | arg_parser.add_argument('-t', '--tx_depth', type=int, default=16, 48 | help='Size of tx queue (default 16).') 49 | arg_parser.add_argument('-T', '--qp_type', type=str, default=IBV_QPT_RC, choices=qp_dict, action=DictAction, 50 | help='The type of QP.') 51 | arg_parser.add_argument('-v', '--version', action='version', version='%(prog)s version : v0.01', 52 | help='Show the version') 53 | arg_parser.add_argument('-x', '--gid_index', type=int, default=1, 54 | help='The index of source gid to communicate with remote peer.') 55 | 56 | # TODO: Check the range and raise an error, for example: 57 | # https://stackoverflow.com/questions/18700634/python-argparse-integer-condition-12 58 | 59 | return vars(arg_parser.parse_args()) 60 | 61 | 62 | parser = ArgsParser() --------------------------------------------------------------------------------