def get_recv_thread(self):
def f():
msg = loads(self.sim2mgr_socket.recv(copy=False).bytes)
self.queue.put(msg)
recv_thread = LoopThread(f, pausable=False)
# recv_thread.daemon = True
recv_thread.name = "recv thread"
return recv_thread
def get_simulator_thread(self):
# spawn a separate thread to run policy, can speed up 1.3x
def populate_job_func():
self._populate_job_queue.get()
with self.trainer.sess.as_default():
for _ in range(self.update_frequency):
self._populate_exp()
th = LoopThread(populate_job_func, pausable=False)
th.name = "SimulatorThread"
return th
def get_recv_thread(self):
def f():
msg = self.sim2exp_socket.recv(copy=False).bytes
msg = loads(msg)
print('{}: received msg'.format(self.agent_name))
try:
self.queue.put_nowait(msg)
except Exception:
logger.info('put queue failed!')
# send response or not?
recv_thread = LoopThread(f, pausable=False)
# recv_thread.daemon = True
recv_thread.name = "recv thread"
return recv_thread
def get_work_thread(self):
def f():
msg = self.queue.get()
sim_name = msg[0]
if msg[1] == SimulatorManager.MSG_TYPE.LOCK and self.locked_sim is None:
self.locked_sim = sim_name
self.mgr2sim_socket.send_multipart(
[sim_name.encode('utf-8'),
dumps('lock')])
time.sleep(0.2)
return
if self.locked_sim is not None:
if sim_name != self.locked_sim:
time.sleep(0.2)
self.queue.put(msg)
return
elif msg[1] == SimulatorManager.MSG_TYPE.UNLOCK:
self.locked_sim = None
self.mgr2sim_socket.send_multipart(
[sim_name.encode('utf-8'),
dumps('unlock')])
time.sleep(0.2)
return
self.cxt_switch(sim_name)
# time.sleep(0.2)
# print(msg[1])
if msg[1] == SimulatorManager.MSG_TYPE.SCREEN:
screen = grab_screen()
self.mgr2sim_socket.send_multipart(
[sim_name.encode('utf-8'),
dumps(screen)])
elif msg[1] == SimulatorManager.MSG_TYPE.CLICK:
# print('need to click')
click(msg[2][0], msg[2][1])
self.mgr2sim_socket.send_multipart(
[sim_name.encode('utf-8'),
dumps('click')])
work_thread = LoopThread(f, pausable=False)
work_thread.name = "work thread"
return work_thread
def setup(self):
super(MultiGPUTrainer, self).setup()
self._training_aux_threads = []
self._training_aux_running = False
self._training_aux_step_counter = itertools.count()
from tensorpack.callbacks.hooks import CallbackToHook
from tensorflow.python.training.monitored_session \
import _HookedSession as HookedSession
for tidx, n in enumerate(self._train_ops_aux):
auxTrainOp = self._train_ops_aux[n]
logger.info("Create aux train ops {}".format(auxTrainOp._name))
if len(auxTrainOp._callbacks) > 0:
auxTrainOp._sess = HookedSession(
self.sess,
hooks=[CallbackToHook(cb) for cb in auxTrainOp._callbacks])
else:
auxTrainOp._sess = self.sess
def f(op=auxTrainOp): # avoid late-binding
try:
op._sess.run([op._train_op
]) # TODO this won't work with StageInput
except RuntimeError: # exited
pass
except tf.errors.CancelledError:
pass
except tf.errors.AbortedError:
pass
# next(self._training_aux_step_counter) # atomic due to GIL
th = LoopThread(f)
th.name = "AsyncLoopThread-{}".format(tidx)
th.pause()
th.start()
logger.info("Start aux thread {}".format(auxTrainOp._name))
self._training_aux_threads.append(th)
