-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathworker_process.py
More file actions
98 lines (84 loc) · 3.4 KB
/
Copy pathworker_process.py
File metadata and controls
98 lines (84 loc) · 3.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
from socket import AF_INET, SOCK_STREAM, socket
from . import RemoteRef, BaseWorker
import pickle
import sys
import torch
from time import sleep
from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
import os
from .config import verbose_worker
# log what is happening
if os.environ['CUDA_VISIBLE_DEVICES'] != '0':
verbose_worker = False
no_response = object()
class RemoteWorker(BaseWorker):
def __init__(self, host, port, secret):
self.socket = socket(AF_INET, SOCK_STREAM)
self.socket.connect((host, int(port)))
self.ofile = self.socket.makefile("wb")
self.ifile = self.socket.makefile("rb")
self._write_pickle(secret)
self.ofile.flush()
self.ref_to_tensor = {}
def run(self):
while True:
(method, *args), dels = self._read_pickle()
for d in dels:
del self.ref_to_tensor[d]
if verbose_worker:
if method == "define_function":
fn, body = args
print(f"method: define_function {fn}\n{body}")
else:
print("method:", method, args)
if method == 'exit':
return
elif method == 'request_value':
ref, = args
result = self.ref_to_tensor[ref.id]
self._write_pickle(result)
self.ofile.flush()
else:
getattr(self, method)(*args)
def send_command(self, func, args, kwargs, results):
def get_tensor(t):
if isinstance(t, RemoteRef):
return self.ref_to_tensor[t.id]
else:
return t
args = tree_map(get_tensor, args)
kwargs = tree_map(get_tensor, kwargs)
result = func(*args, **kwargs)
flat_results, _ = tree_flatten(result)
real_results = [e for e in flat_results if isinstance(e, torch.Tensor)]
for real, ref in zip(real_results, results):
self.ref_to_tensor[ref.id] = real
def send_value(self, ref: RemoteRef, value: torch.Tensor):
self.ref_to_tensor[ref.id] = value
def del_value(self, ref: RemoteRef):
del self.ref_to_tensor[ref.id]
def _write_pickle(self, obj):
b = pickle.dumps(obj)
sz = len(b).to_bytes(8, 'little')
self.ofile.write(sz)
self.ofile.write(b)
def _read_pickle(self):
sz = int.from_bytes(self.ifile.read(8), 'little')
return pickle.loads(self.ifile.read(sz))
def create_process_group(self, rank, world_size, pg_ref):
torch.distributed.init_process_group('nccl', init_method='tcp://127.0.0.1:12350', rank=rank, world_size=world_size)
self.ref_to_tensor[pg_ref.id] = None
def create_process_subgroup(self, orig_pg, participating_ranks, pg):
pg_obj = self.ref_to_tensor[orig_pg.id]
assert pg_obj is None, "subgroup must be created from default group..."
r = torch.distributed.new_group(ranks=participating_ranks, backend='nccl')
if pg is not None:
self.ref_to_tensor[pg.id] = r
def all_reduce_coalesced(self, refs: 'List[int]', pg_ref: RemoteRef):
pg = self.ref_to_tensor[pg_ref.id]
ts = [self.ref_to_tensor[r] for r in refs]
torch.distributed.all_reduce_coalesced(ts, group=pg)
if __name__ == "__main__":
_, host, port, secret = sys.argv
w = RemoteWorker(host, port, secret)
w.run()