-
Notifications
You must be signed in to change notification settings - Fork 0
/
batch_main.py
81 lines (65 loc) · 2.4 KB
/
batch_main.py
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
'''build a tmux session to maintain
each python command_to_run in a window'''
import time
start_time = time.time()
import argparse
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--env-names', type=str, nargs='*',)
parser.add_argument('--cards', type=str, nargs='*',)
parser.add_argument('--agent-name', type=str,)
parser.add_argument('--command', type=str,)
args = parser.parse_args()
command_to_run = 'source activate Mega-Agent-2 && CUDA_VISIBLE_DEVICES=CARD python main.py --env-name GAME '
game_append = 'NoFrameskip-v4'
if args.agent_name is None:
args.agent_name = input('# ACTION REQUIRED: Name the agent:')
print('# INFO: Name of the agent: {}'.format(args.agent_name))
if args.command is None:
args.command = input('# ACTION REQUIRED: Provide command:')
print('# INFO: Command is: {}'.format(args.command))
command_to_run += args.command
assert len(args.cards)==len(args.env_names)
import libtmux
server = libtmux.Server()
try:
server.kill_session(args.agent_name)
print('# INFO: Previous session {} is killed'.format(args.agent_name))
except Exception as e:
print('# INFO: No previous session is killed')
session = server.new_session(args.agent_name)
for i in range(len(args.cards)):
'''build each command_to_run a window to maintain it'''
pane = session.new_window(
attach = False,
window_name = '{}'.format(
args.env_names[i],
)
).list_panes()[0]
'''send command to run command_to_run in the window's pane'''
pane.send_keys(
command_to_run.replace(
'CARD',str(args.cards[i])
).replace(
'GAME',(str(args.env_names[i])+game_append)
)
)
'''build a signal handler to kill all windows
session when the main thread is killed'''
import signal
import sys
def signal_handler(signal, frame):
print()
print('# WARNING: You pressed Ctrl+C! Now, kill the session {} after running for {:.2f} hours'.format(
args.agent_name,
(time.time()-start_time)/60.0/60.0,
))
session.kill_session()
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
while True:
input('# INFO: Up for {:.2f} hours. Running {} agents over {} on cards {}. Kill this thread with Ctrl+C to kill the session.'.format(
(time.time()-start_time)/60.0/60.0,
args.agent_name,
args.env_names,
args.cards,
))