feature: cumulative worker id without entry server

handyrl/worker.py

@@ -84,8 +84,8 @@ def run(self):
                 send_recv(self.conn, ('result', result))
 
 
-def make_worker_args(args, n_ga, gaid, wid, conn):
-    return args, conn, wid * n_ga + gaid
+def make_worker_args(args, base_wid, wid, conn):
+    return args, conn, base_wid + wid
 
 
 def open_worker(args, conn, wid):
@@ -94,23 +94,20 @@ def open_worker(args, conn, wid):
 
 
 class Gather(QueueCommunicator):
-    def __init__(self, args, conn, gaid):
-        print('started gather %d' % gaid)
+    def __init__(self, args, conn, gather_id, base_worker_id, num_workers):
+        print('started gather %d' % gather_id)
         super().__init__()
-        self.gather_id = gaid
+        self.gather_id = gather_id
         self.server_conn = conn
         self.args_queue = deque([])
         self.data_map = {'model': {}}
         self.result_send_map = {}
         self.result_send_cnt = 0
 
-        n_pro, n_ga = args['worker']['num_parallel'], args['worker']['num_gathers']
-
-        num_workers_per_gather = (n_pro // n_ga) + int(gaid < n_pro % n_ga)
         worker_conns = open_multiprocessing_connections(
-            num_workers_per_gather,
+            num_workers,
             open_worker,
-            functools.partial(make_worker_args, args, n_ga, gaid)
+            functools.partial(make_worker_args, args, base_worker_id)
         )
 
         for conn in worker_conns:
@@ -160,9 +157,25 @@ def run(self):
                     self.result_send_cnt = 0
 
 
-def gather_loop(args, conn, gaid):
+def gather_loop(args, conn, gather_id):
+    n_pro, n_ga = args['worker']['num_parallel'], args['worker']['num_gathers']
+    n_pro_w = (n_pro // n_ga) + int(gather_id < n_pro % n_ga)
+    args['worker']['num_parallel_per_gather'] = n_pro_w
+    base_worker_id = 0
+
+    if conn is None:
+        # entry
+        conn = connect_socket_connection(args['worker']['server_address'], 9998)
+        conn.send(args['worker'])
+        args = conn.recv()
+
+        if gather_id == 0:  # call once at every machine
+             print(args)
+             prepare_env(args['env'])
+        base_worker_id = args['worker'].get('base_worker_id', 0)
+
     try:
-        gather = Gather(args, conn, gaid)
+        gather = Gather(args, conn, gather_id, base_worker_id, n_pro_w)
         gather.run()
     finally:
         gather.shutdown()
@@ -188,47 +201,31 @@ class WorkerServer(QueueCommunicator):
     def __init__(self, args):
         super().__init__()
         self.args = args
+        self.total_worker_count = 0
 
     def run(self):
         # prepare listening connections
-        def entry_server(port):
-            print('started entry server %d' % port)
+        def worker_server(port):
+            print('started worker server %d' % port)
             conn_acceptor = accept_socket_connections(port=port, timeout=0.3)
             while not self.shutdown_flag:
                 conn = next(conn_acceptor)
                 if conn is not None:
                     worker_args = conn.recv()
                     print('accepted connection from %s!' % worker_args['address'])
+                    worker_args['base_worker_id'] = self.total_worker_count
+                    self.total_worker_count += worker_args['num_parallel_per_gather']
                     args = copy.deepcopy(self.args)
                     args['worker'] = worker_args
                     conn.send(args)
-                    conn.close()
-            print('finished entry server')
-
-        def worker_server(port):
-            print('started worker server %d' % port)
-            conn_acceptor = accept_socket_connections(port=port, timeout=0.3)
-            while not self.shutdown_flag:
-                conn = next(conn_acceptor)
-                if conn is not None:
                     self.add_connection(conn)
             print('finished worker server')
 
         # use thread list of super class
-        self.threads.append(threading.Thread(target=entry_server, args=(9999,)))
         self.threads.append(threading.Thread(target=worker_server, args=(9998,)))
-        self.threads[-2].start()
         self.threads[-1].start()
 
 
-def entry(worker_args):
-    conn = connect_socket_connection(worker_args['server_address'], 9999)
-    conn.send(worker_args)
-    args = conn.recv()
-    conn.close()
-    return args
-
-
 class RemoteWorkerCluster:
     def __init__(self, args):
         args['address'] = gethostname()
@@ -238,18 +235,12 @@ def __init__(self, args):
         self.args = args
 
     def run(self):
-        args = entry(self.args)
-        print(args)
-        prepare_env(args['env'])
-
         # open worker
         process = []
         try:
             for i in range(self.args['num_gathers']):
-                conn = connect_socket_connection(self.args['server_address'], 9998)
-                p = mp.Process(target=gather_loop, args=(args, conn, i))
+                p = mp.Process(target=gather_loop, args=({'worker': self.args}, None, i))
                 p.start()
-                conn.close()
                 process.append(p)
             while True:
                 time.sleep(100)