def test_set_status(self):
with api.app.app_context():
with self.assertRaises(ModelNotFound):
set_model(test_version)
save_model('Hello', test_version)
set_model(test_version)
self.assertEqual(get_status()[-1]['current_version'], test_version)
def test_get_status(self):
with api.app.app_context():
save_model('hi', version=test_version)
save_model('hi')
set_model(test_version + 1)
set_model(test_version)
response = self.app.get('/status')
data = json.loads(response.data.decode('utf-8'))
self.assertEqual(len(data), 2)
self.assertEqual(data[-1]['current_version'], test_version)
def test_api_set_status(self):
with api.app.app_context():
save_model('hi', version=test_version)
save_model('hi')
self.assertEqual(get_current_model(),
MODEL_ROOT_NAME + str(test_version + 1))
self.app.put('/status',
headers={'content-type': 'application/json'},
data=json.dumps({'version': test_version}))
with api.app.app_context():
self.assertEqual(get_current_model(),
MODEL_ROOT_NAME + str(test_version))
def test_get_current_model(self):
with api.app.app_context():
with self.assertRaises(NoModelsFound):
get_current_model()
save_model('Hello', test_version)
save_model('Hello')
# no status set yet, should get latest model
self.assertEqual(get_current_model(),
MODEL_ROOT_NAME + str(test_version + 1))
set_model(test_version)
# status now set to previous, should get that model
self.assertEqual(get_current_model(),
MODEL_ROOT_NAME + str(test_version))
def run_retrain(source_filename, hyperparams=None, test=False):
with app.app_context(): # This is used since the run_retrain requires app context
target, version_number = next_model_name()
logger = get_logger(not test)
if test:
logger.info('Starting test evaluation from source json dump {}'.format(source_filename, target))
else:
logger.info('Starting full retrain from source json dump {} to model file {}'.format(source_filename, target))
sync_log_to_s3(logger)
try:
source = load_train_file_into_memory(source_filename)
except Exception as ex:
if ex.response['Error']['Code'] == '404':
logger.error(ex)
sync_log_to_s3(logger)
abort(message='Retrain file not found in twc-input s3 bucket')
try:
_, _, _, model = train_model_from_json(source, hyperparams=hyperparams, test=test)
if test:
run_retrain(source_filename, hyperparams, False)
else:
pickle.dump(model, open(target, 'wb'))
log_filename = '{}_retrain_{}.log'.format(source_filename, target)
logger.info('Uploading model file [{}] and log receipt [{}] to s3'.format(target, log_filename))
sync_log_to_s3(logger)
log_file = open('retrain_output.log', 'r').read()
save_model(target, 'retrain_output.log', log_filename)
set_model(version_number)
clear_log_file_from_s3()
return log_file
except Exception as ex:
logger.error(ex)
log_file = open('retrain_output.log', 'r').read()
abort(message=log_file)
def test_overwrite_model(self):
with api.app.app_context():
save_model('Hello', test_version)
with self.assertRaises(OverwriteFailure):
save_model('Hello', test_version)
def test_save_models(self):
with api.app.app_context():
save_model('Hello', test_version)
models = get_models()
self.assertTrue(test_version in models)
def test_get_models(self):
with api.app.app_context():
save_model('hi', version=test_version)
save_model('hi')
response = self.app.get('/models')
self.assertEqual(len(json.loads(response.data.decode('utf-8'))), 2)