Skip to content
Snippets Groups Projects
Commit ccb3c24f authored by Rishi Sharma's avatar Rishi Sharma
Browse files

Add compression

parent 5ab284a8
No related branches found
No related tags found
No related merge requests found
import json import json
import logging import logging
import lzma
import pickle import pickle
from collections import deque from collections import deque
...@@ -38,7 +39,9 @@ class TCP(Communication): ...@@ -38,7 +39,9 @@ class TCP(Communication):
port = rank + 20000 port = rank + 20000
return "tcp://{}:{}".format(machine_addr, port) return "tcp://{}:{}".format(machine_addr, port)
def __init__(self, rank, machine_id, mapping, total_procs, addresses_filepath): def __init__(
self, rank, machine_id, mapping, total_procs, addresses_filepath, compress=False
):
""" """
Constructor Constructor
...@@ -72,6 +75,7 @@ class TCP(Communication): ...@@ -72,6 +75,7 @@ class TCP(Communication):
self.router.setsockopt(zmq.IDENTITY, self.identity) self.router.setsockopt(zmq.IDENTITY, self.identity)
self.router.bind(self.addr(rank, machine_id)) self.router.bind(self.addr(rank, machine_id))
self.sent_disconnections = False self.sent_disconnections = False
self.compress = compress
self.peer_deque = deque() self.peer_deque = deque()
self.peer_sockets = dict() self.peer_sockets = dict()
...@@ -99,7 +103,13 @@ class TCP(Communication): ...@@ -99,7 +103,13 @@ class TCP(Communication):
Encoded data Encoded data
""" """
return pickle.dumps(data) if self.compress:
compressor = lzma.LZMACompressor()
output = compressor.compress(pickle.dumps(data)) + compressor.flush()
else:
output = pickle.dumps(data)
return output
def decrypt(self, sender, data): def decrypt(self, sender, data):
""" """
...@@ -119,7 +129,10 @@ class TCP(Communication): ...@@ -119,7 +129,10 @@ class TCP(Communication):
""" """
sender = int(sender.decode()) sender = int(sender.decode())
data = pickle.loads(data) if self.compress:
data = pickle.loads(lzma.decompress(data))
else:
data = pickle.loads(data)
return sender, data return sender, data
def connect_neighbors(self, neighbors): def connect_neighbors(self, neighbors):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment