/
fabfile.py
158 lines (136 loc) · 5.86 KB
/
fabfile.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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
import os
from datetime import datetime
from fabric.api import env, execute
from fabric.utils import puts
from fabric.operations import run, put, local
from fabric import colors
import fabric.network
import fabric.state
YAML_AVAILABLE = True
try:
import yaml
except ImportError:
YAML_AVAILABLE = False
JSON_AVAILABLE = True
try:
import simplejson as json
except ImportError:
try:
import json
except ImportError:
JSON_AVAILABLE = False
################################
# ENVIRONMENTS #
################################
def _load_config(**kwargs):
"""Find and parse server config file.
If `config` keyword argument wasn't set look for default
'server_config.yaml' or 'server_config.json' file.
"""
config, ext = os.path.splitext(kwargs.get('config',
'server_config.yaml' if os.path.exists('config.yaml') else 'config.json'))
if not os.path.exists(config + ext):
print colors.red('Error. "%s" file not found.' % (config + ext))
return {}
if YAML_AVAILABLE and ext == '.yaml':
loader = yaml
elif JSON_AVAILABLE and ext =='.json':
loader = json
else:
print colors.red('Parser package not available')
return {}
# Open file and deserialize settings.
with open(config + ext) as config_file:
return loader.load(config_file)
def s(*args, **kwargs):
"""Set destination servers or server groups by comma delimited list of names"""
# Load config
servers = _load_config(**kwargs)
# If no arguments were recieved, print a message with a list of available configs.
if not args:
print 'No server name given. Available configs:'
for key in servers:
print colors.green('\t%s' % key)
# Create `group` - a dictionary, containing copies of configs for selected servers. Server hosts
# are used as dictionary keys, which allows us to connect current command destination host with
# the correct config. This is important, because somewhere along the way fabric messes up the
# hosts order, so simple list index incrementation won't suffice.
env.group = {}
# For each given server name
for name in args:
# Recursive function call to retrieve all server records. If `name` is a group(e.g. `all`)
# - get it's members, iterate through them and create `group`
# record. Else, get fields from `name` server record.
# If requested server is not in the settings dictionary output error message and list all
# available servers.
_build_group(name, servers)
# Copy server hosts from `env.group` keys - this gives us a complete list of unique hosts to
# operate on. No host is added twice, so we can safely add overlaping groups. Each added host is
# guaranteed to have a config record in `env.group`.
env.hosts = env.group.keys()
def _build_group(name, servers):
"""Recursively walk through servers dictionary and search for all server records."""
# We're going to reference server a lot, so we'd better store it.
server = servers.get(name, None)
# If `name` exists in servers dictionary we
if server:
# check whether it's a group by looking for `members`
if isinstance(server, list):
if fabric.state.output['debug']:
puts("%s is a group, getting members" % name)
for item in server:
# and call this function for each of them.
_build_group(item, servers)
# When, finally, we dig through to the standalone server records, we retrieve
# configs and store them in `env.group`
else:
if fabric.state.output['debug']:
puts("%s is a server, filling up env.group" % name)
env.group[server['host']] = server
else:
print colors.red('Error. "%s" config not found. Run `fab s` to list all available configs' % name)
def _setup(task):
"""
Copies server config settings from `env.group` dictionary to env variable.
This way, tasks have easier access to server-specific variables:
`env.owner` instead of `env.group[env.host]['owner']`
"""
def task_with_setup(*args, **kwargs):
# If `s:server` was run before the current command - then we should copy values to
# `env`. Otherwise, hosts were passed through command line with `fab -H host1,host2
# command` and we skip.
if env.get("group", None):
for key,val in env.group[env.host].items():
setattr(env, key, val)
if fabric.state.output['debug']:
puts("[env] %s : %s" % (key, val))
task(*args, **kwargs)
# Don't keep host connections open, disconnect from each host after each task.
# Function will be available in fabric 1.0 release.
# fabric.network.disconnect_all()
return task_with_setup
#############################
# TASKS #
#############################
current_time_stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@_setup
def execute_methods():
for i in range(0, len(env.gpus)):
remote_pkg_path = os.path.join('~',env['package-names'][i])
run('rm -rf {0}'.format(remote_pkg_path))
run('mkdir -p {0}'.format(remote_pkg_path))
put(os.path.join(env['local-package-name'],'*'), remote_pkg_path)
run('bash {0} {1} {2} {3} {4} {5} {6}'.format(os.path.join(remote_pkg_path, env['profiling-script']),
env.gpus[i].replace(' ','_'),
current_time_stamp,
env['upload-google-drive'],
env['upload-dropbox'],
env.prefixes[i],
env['package-names'][i]))
def generate_data():
if env.host == 'localhost':
local('bash {0} {1}'.format(env['generate-script'], current_time_stamp))
def start_profiling():
execute(execute_methods)
execute(generate_data)
print colors.green('Profiling has been finished!')