-
Notifications
You must be signed in to change notification settings - Fork 23
/
setup.py
222 lines (186 loc) · 8.64 KB
/
setup.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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
# After setting up the config.json, lambda-function.py, and trainer-script.sh files,
# run this script to create the necessary Lambda function. Execution of the Lambda
# function will be triggered manually by you from the AWS CLI or web console, a separate
# program of yours via the AWS SDK, an IoT button, or whatever other trigger you make use of
# for firing off your algorithm's training session.
import boto3
import json
import logging
import os
import zipfile
def lambda_packer():
"""
To create the Lambda function we need to pack up a ZIP file to upload. Returns true if the ZIP archive is created,
False is something goes wrong. If it succeeds, the second item in the returned tuple object is the 'lambdapack'
filepath, or the relative filepath to the Lambda-ready ZIP file.
:return:
"""
try:
# Create a ZIP archive, with compression. The files in `config/` and `src/` are the only ones we need to include
# for the Lambda to function, but in case you have custom needs you can add the necessary directories under this
# dict.
directories_to_include = ['config', 'src']
destination_filepath = 'target/Parris-v1-Lambda.zip'
# Create the target directory if it doesn't exist.
if not os.path.exists('target'):
try:
os.makedirs('target')
except Exception as e:
raise Exception('Tried to create target/ dir but failed. {}'.format(e))
with zipfile.ZipFile(destination_filepath, mode='w') as lambdapack:
# Crawl directories, and write their files into the ZIP with the structure of foldername/filename.
# This should maintain subdirectory heirarchies.
for dirname in directories_to_include:
for root, dir, filename in os.walk(dirname):
for individual_file in filename:
composite_filename = os.path.join(root, individual_file)
logging.debug('Adding: {}'.format(composite_filename))
# AWS Lambda functions are expected to have the handler at the top level of the package.
# If you have other scripts like lambda-function.py that need to be in the root of the package,
# include them in the list below and they'll be put at the root of the ZIP archive.
if individual_file in ['lambda-function.py']:
lambdapack.write(filename=composite_filename, arcname=individual_file)
else:
lambdapack.write(filename=composite_filename, arcname=composite_filename)
logging.info('Packed lambdapack to {}'.format(destination_filepath))
return [True, destination_filepath]
except Exception as e:
msg = 'lambda_packer failure: {}'.format(e)
logging.error(msg)
return [False, msg]
def lambda_creation(lambda_config={}, lambdapack='', dryrun=False):
"""
Create the lambda function, return the response, give the AWS CLI command to execute it.
:return:
"""
try:
# Make a Lambda Boto3 client, upload lambdapack, return successful ARN to prove success.
client_lambda = boto3.client('lambda')
try:
logging.warning('Uploading Lambdapack from {}'.format(lambdapack))
creation_response = client_lambda.create_function(
FunctionName='Parris-v1-Lambda',
Runtime='python3.6',
Role=lambda_config['lambda-role-arn'],
Handler='lambda-function.lambda_handler',
Code={
'ZipFile': open(lambdapack, mode='rb').read()
},
Description='Lambda function for Parris, the ML training automation tool.',
Timeout=30,
MemorySize=128,
Environment={
'Variables': {
's3_training_bucket': lambda_config.get('s3-training-bucket', '')
}
},
Tags={
'Name': 'Parris-v1-Lambda'
}
)
except Exception as update_err:
# If an error was thrown during the creation of the Lambda function, we want to intercept this and check if
# it was because the function already exists. If it does, try an update to it. Otherwise pass it along.
# Normally we would do an 'except ResourceConflictException:' block here, but for whatever reason this class
# is not available in botocore or boto3 to use, so this'll have to do.
if 'ResourceConflictException' in str(update_err):
logging.error('lambda_creation encountered ResourceConflictException - attempting function update.')
if dryrun:
logging.warning('lambda_creation is running under DryRun mode for the update function.')
creation_response = client_lambda.update_function_code(
FunctionName='Parris-v1-Lambda',
ZipFile=open(lambdapack, mode='rb').read(),
DryRun=dryrun
)
else:
raise Exception(update_err)
# Report success with the function's ARN!
try:
lambda_arn = creation_response['FunctionArn']
except Exception as arn_err:
raise Exception(
'lambda_creation failure: Lambda ARN not pulled from response: {} \nResponse contents: {}'
.format(arn_err, creation_response)
)
logging.warning('Successfully created function: \n{}'.format(lambda_arn))
return [True, lambda_arn]
except Exception as e:
msg = 'lambda_creation failure: {}'.format(e)
logging.error(msg)
return [False, msg]
def parse_config():
"""
Parse and return the config. This'll be a lot shorter than the _test function.
:return:
"""
try:
config = json.load(open('config/lambda-config.json'))
return config
except Exception as e:
msg = 'parse_config failure: {}'.format(e)
logging.error(msg)
return False
def _test_lambda_packer():
"""
Try to pack up a Lambdapack with the current config.
:return:
"""
try:
lambda_packer()
return [True, '']
except Exception as e:
msg = '_test_lambda_packer failure: {}'.format(e)
logging.error(msg)
return [False, msg]
def _test_lambda_creation():
"""
Parse and return the config. This'll be a lot shorter than the _test function.
:return:
"""
try:
lambda_config = parse_config()
lambdapack_successfail, lambdapack_filepath = lambda_packer()
lambda_creation(lambda_config, lambdapack=lambdapack_filepath, dryrun=True)
return [True, '']
except Exception as e:
msg = '_test_lambda_creation failure: {}'.format(e)
logging.error(msg)
return [False, msg]
def _test_parse_config():
"""
Try parsing the JSON config. Check for file existence, valid parsing, and of course
that valid inputs are configured. (Numerical inputs for time and cost, for example.)
:return:
"""
try:
lambda_config = parse_config()
logging.debug('Config parsed, training-job-name set to: {}'.format(lambda_config['lambda-role-arn']))
return [True, '']
except Exception as e:
msg = '_test_parse_config failure: {}'.format(e)
logging.error(msg)
return [False, msg]
if __name__ == '__main__':
"""
Parse the config, craft a CloudFormation template, create the Lambda function.
"""
testresult = []
testresult.append(_test_parse_config())
testresult.append(_test_lambda_packer())
# There is no dryrun option for creating the function so this isn't terribly helpful for us.
# testresult.append(_test_lambda_creation())
for testresult, testmsg in testresult:
if not testresult:
raise Exception('Tests did not pass: {}'.format(testmsg))
logging.warning('All setup.py tests passed!')
# If the tests have all passed, go ahead with function creation.
# Lambdapack is your Lambda-ready ZIP file. You don't have to use the lambda_packer function if you already have one
# made up, but this is likely the option you want.
try:
lambda_config = parse_config()
lambdapack_successfail, lambdapack_filepath = lambda_packer()
lambda_creation(lambda_config=lambda_config, lambdapack=lambdapack_filepath)
except Exception as e:
msg = '_test_parse_config failure: {}'.format(e)
logging.error(msg)
logging.info('setup.py finished.')