def start_mxnet_model_server():
serving_env = env.ServingEnv()
is_multi_model = True
if serving_env.module_name is None:
logging.info("Starting MXNet server in algorithm mode.")
_start_model_server(is_multi_model, ALGO_HANDLER_SERVICE)
else:
logging.info("Staring MXNet Model Server with user module.")
# Install user module from s3 to import
modules.import_module(serving_env.module_dir, serving_env.module_name)
_start_model_server(is_multi_model, USER_HANDLER_SERVICE)
def test_import_module_with_local_script(user_module, user_module_name, tmpdir):
tmp_code_dir = str(tmpdir)
user_module.create_tmp_dir_with_files(tmp_code_dir)
module = modules.import_module(tmp_code_dir, user_module_name, cache=False)
assert module.validate()
def test_import_module(user_module_name):
user_module = test.UserModule(USER_SCRIPT).add_file(SETUP).upload()
module = modules.import_module(user_module.url,
user_module_name,
cache=False)
assert module.validate()
def test_import_module_with_s3_script(user_module, user_module_name):
user_module.upload()
module = modules.import_module(user_module.url,
user_module_name,
cache=False)
assert module.validate()
def test_import_module_with_s3_script_with_additional_files(user_module_name):
user_module = test.UserModule(USER_SCRIPT_WITH_ADDITIONAL_FILE).add_file(
ADDITIONAL_FILE).upload()
module = modules.import_module(user_module.url,
user_module_name,
cache=False)
assert module.validate()
def test_import_module_with_s3_script_with_requirements(
user_module, user_module_name):
user_module = user_module.add_file(REQUIREMENTS_FILE).upload()
module = modules.import_module(user_module.url,
user_module_name,
cache=False)
assert module.say() == REQUIREMENTS_TXT_ASSERT_STR
def main(environ, start_response):
global app
if app is None:
serving_env = env.ServingEnv()
user_module = modules.import_module(serving_env.module_dir,
serving_env.module_name)
user_module_transformer = _user_module_transformer(user_module)
user_module_transformer.initialize()
app = worker.Worker(transform_fn=user_module_transformer.transform,
module_name=serving_env.module_name)
return app(environ, start_response)
def test_import_module_with_local_tar_via_download_and_extract(
user_module, user_module_name):
user_module = user_module.add_file(REQUIREMENTS_FILE)
tar_name = user_module.create_tar()
module = modules.import_module(tar_name,
name=user_module_name,
cache=False)
assert module.say() == REQUIREMENTS_TXT_ASSERT_STR
os.remove(tar_name)
def main(environ, start_response):
global app
if app is None:
serving_env = env.ServingEnv()
_update_mxnet_env_vars()
user_module = modules.import_module(serving_env.module_dir, serving_env.module_name)
user_transformer = _user_module_transformer(user_module, serving_env.model_dir)
app = worker.Worker(transform_fn=user_transformer.transform,
module_name=serving_env.module_name)
return app(environ, start_response)
def main(environ, start_response):
global app
if app is None:
serving_env = env.ServingEnv()
logger.setLevel(serving_env.log_level)
user_module = modules.import_module(serving_env.module_dir, serving_env.module_name)
user_module_transformer = _user_module_transformer(user_module)
user_module_transformer.initialize()
app = worker.Worker(transform_fn=user_module_transformer.transform,
module_name=serving_env.module_name,
healthcheck_fn=default_healthcheck_fn)
return app(environ, start_response)
def framework_training_fn():
training_env = sagemaker_containers.training_env()
mod = modules.import_module(training_env.module_dir, training_env.module_name, False)
model = mod.train(**functions.matching_args(mod.train, training_env))
if model:
if hasattr(mod, 'save'):
mod.save(model, training_env.model_dir)
else:
model_file = os.path.join(training_env.model_dir, 'saved_model')
model.save(model_file)
def train():
training_env = sagemaker_containers.training_env()
script = modules.import_module(training_env.module_dir,
training_env.module_name)
model = script.train(**functions.matching_args(script.train, training_env))
if model:
if hasattr(script, "save"):
script.save(model, training_env.model_dir)
else:
model_file = os.path.join(training_env.model_dir, "saved_model")
model.save(model_file)
def import_module(module_name, module_dir):
try: # if module_name already exists, use the existing one
user_module = importlib.import_module(module_name)
except ImportError: # if the module has not been loaded, 'modules' downloads and installs it.
user_module = modules.import_module(module_dir, module_name)
except Exception: # this shouldn't happen
logger.info("Encountered an unexpected error.")
raise
user_module_transformer = _user_module_transformer(user_module)
user_module_transformer.initialize()
return user_module_transformer
def test_import_module_with_s3_script_with_requirements(user_module_name):
user_module = test.UserModule(USER_SCRIPT_WITH_REQUIREMENTS).add_file(
REQUIREMENTS_FILE).upload()
module = modules.import_module(user_module.url,
user_module_name,
cache=False)
assert module.say() == """
____ __ __ _.............
/ ___| __ _ __ _ ___| \/ | __ _| | _____ _ __.
\___ \ / _` |/ _` |/ _ \ |\/| |/ _` | |/ / _ \ '__|
___) | (_| | (_| | __/ | | | (_| | < __/ |...
|____/ \__,_|\__, |\___|_| |_|\__,_|_|\_\___|_|...
|___/.................................
""".replace('.', ' ').strip()
def start_model_server():
serving_env = env.ServingEnv()
is_multi_model = True
modules.import_module(serving_env.module_dir, serving_env.module_name)
_start_model_server(is_multi_model, HANDLER_SERVICE)
def test_import_module(user_module, user_module_name):
user_module.upload()
module = modules.import_module(user_module.url, user_module_name)
assert module.validate()
def test_import_module_with_s3_script_with_error(user_module_name):
user_module = test.UserModule(USER_SCRIPT_WITH_ERROR).add_file(
SETUP_FILE).upload()
with pytest.raises(errors.ImportModuleError):
modules.import_module(user_module.url, user_module_name, cache=False)
