The current implementation of the P2P data streaming in the Connection class allows nodes to send potentially large amounts of data without sufficient structure or limitations. This can lead to inefficiencies, unintended data transmission, or even exploitation by malicious nodes. Enhancements are needed to impose stricter control over the data streaming process, ensuring that only expected and valid data is transmitted and received.
Proposed changes
- Data expectation framework: Introduce a mechanism to define expected data structures for
connection.py from node instances, especially when expecting the data (e.g. when expecting a module or tensors from a worker/user).
- Perhaps with the use of a callback function from
Conenction to SmartNode when receiving tagged data to check that we are expecting or can handle it.
- Chunk size enforcement: Reduce chunk size and maximum data to be received. Only increase these limits momentarily when expecting larger data like modules, tensors, or state updates.
- Rate limiter attempted in
smart_node.py could be modularized and set up to work in connection.py as well.
- incorporate a blacklist that is stored to disk (perhaps logs/dht_state.json)
The changes will primarily impact the Connection class in p2p/connection.py, as well as the handle_data methods of nodes and roles (i.e. smart_node.py, torch_node.py, validator.py, user.py, and worker.py)
The current implementation of the P2P data streaming in the Connection class allows nodes to send potentially large amounts of data without sufficient structure or limitations. This can lead to inefficiencies, unintended data transmission, or even exploitation by malicious nodes. Enhancements are needed to impose stricter control over the data streaming process, ensuring that only expected and valid data is transmitted and received.
Proposed changes
connection.pyfrom node instances, especially when expecting the data (e.g. when expecting a module or tensors from a worker/user).ConenctiontoSmartNodewhen receiving tagged data to check that we are expecting or can handle it.smart_node.pycould be modularized and set up to work inconnection.pyas well.The changes will primarily impact the Connection class in p2p/connection.py, as well as the handle_data methods of nodes and roles (i.e. smart_node.py, torch_node.py, validator.py, user.py, and worker.py)