def uninstall_requirements_file():
requirements_data = "pyfiglet"
yield test.File("requirements.txt", requirements_data)
try:
subprocess.check_call(
shlex.split("pip uninstall -y --quiet %s" % requirements_data))
except subprocess.CalledProcessError:
pass
def test_script_mode_client_import_error():
channel = test.Channel.create(name='training')
requirements_file = test.File('requirements.txt', '42/0')
user_script = test.File(name='user_script.py', data='42/0')
module = test.UserModule(user_script).add_file(requirements_file).upload()
hyperparameters = dict(sagemaker_program='user_script.py')
test.prepare(user_module=module,
hyperparameters=hyperparameters,
channels=[channel])
with pytest.raises(errors.InstallModuleError) as e:
framework_training_with_script_mode_fn()
message = str(e.value)
assert 'InstallModuleError:' in message
def test_parameter_server():
module = test.UserModule(test.File(name='user_script.py', data=PARAMETER_SERVER_SCRIPT))
hyperparameters = dict(sagemaker_program='user_script.py')
test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[test.Channel.create(name='training')])
training_env = sagemaker_containers.training_env()
process = entry_point.run(training_env.module_dir, training_env.user_entry_point,
training_env.to_cmd_args(), training_env.to_env_vars(), wait=False)
# confirm the ps process is still hanging
assert process.poll() is None
process.kill()
def test_script_mode_client_import_error(training_fn, capture_error):
channel = test.Channel.create(name='training')
requirements_file = test.File('requirements.txt', '42/0')
user_script = test.File(name='user_script', data='42/0')
module = test.UserModule(user_script).add_file(setup_file).add_file(requirements_file).upload()
hyperparameters = dict(sagemaker_program='user_script')
test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
with pytest.raises(errors.InstallModuleError) as e:
training_fn(capture_error)
message = str(e.value)
assert 'InstallModuleError:' in message
if capture_error:
assert "Invalid requirement: \'42/0\'" in message
assert "It looks like a path. File \'42/0\' does not exist." in message
def test_script_mode_client_import_error():
channel = test.Channel.create(name="training")
requirements_file = test.File("requirements.txt", "invalid/module")
user_script = test.File(name="user_script", data="invalid/module")
module = test.UserModule(user_script).add_file(setup_file).add_file(requirements_file).upload()
hyperparameters = dict(sagemaker_program="user_script")
test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
with pytest.raises(errors.InstallModuleError) as e:
framework_training_with_script_mode_fn(capture_error=True)
message = str(e.value)
assert "InstallModuleError:" in message
# fmt: off
assert "Invalid requirement" in message
assert "It looks like a path" in message
def test_script_mode_client_error():
channel = test.Channel.create(name="training")
module = test.UserModule(test.File(name="user_script.py", data=USER_MODE_SCRIPT_WITH_ERROR))
hyperparameters = dict(sagemaker_program="user_script.py")
test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
with pytest.raises(errors.ExecuteUserScriptError) as e:
framework_training_with_script_mode_fn(capture_error=True)
message = str(e.value)
assert "ExecuteUserScriptError" in message
assert "ZeroDivisionError" in message
def test_script_mode_client_error(training_fn, capture_error):
channel = test.Channel.create(name='training')
module = test.UserModule(test.File(name='user_script.py', data=USER_MODE_SCRIPT_WITH_ERROR))
hyperparameters = dict(sagemaker_program='user_script.py')
test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
with pytest.raises(errors.ExecuteUserScriptError) as e:
training_fn(capture_error)
message = str(e.value)
assert 'ExecuteUserScriptError' in message
if capture_error:
assert 'ZeroDivisionError' in message
def test_training_framework_network_isolation(user_script, capture_error):
with pytest.raises(ImportError):
importlib.import_module(modules.DEFAULT_MODULE_NAME)
channel_train = test.Channel.create(name="training")
channel_code = test.Channel.create(name="code")
features = [1, 2, 3, 4]
labels = [0, 1, 0, 1]
np.savez(os.path.join(channel_train.path, "training_data"),
features=features,
labels=labels)
with open(os.path.join(channel_code.path, "user_script.py"), "w") as f:
f.write(user_script)
f.close()
module = test.UserModule(test.File(
name="user_script.py", data="")) # dummy module for hyperparameters
submit_dir = env.input_dir + "/data/code"
hyperparameters = dict(
training_data_file="training_data.npz",
sagemaker_program="user_script.py",
sagemaker_submit_directory=submit_dir,
epochs=10,
batch_size=64,
optimizer="Adam",
)
test.prepare(
user_module=module,
hyperparameters=hyperparameters,
channels=[channel_train, channel_code],
local=True,
)
assert execute_an_wrap_exit(framework_training_fn) == trainer.SUCCESS_CODE
model_path = os.path.join(env.model_dir, "saved_model")
model = fake_ml_framework.Model.load(model_path)
assert model.epochs == 10
assert model.batch_size == 64
assert model.optimizer == "Adam"
def test_script_mode_local_directory(user_script, training_fn, capture_error,
tmpdir):
channel = test.Channel.create(name="training")
features = [1, 2, 3, 4]
labels = [0, 1, 0, 1]
np.savez(os.path.join(channel.path, "training_data"),
features=features,
labels=labels)
tmp_code_dir = str(tmpdir)
module = test.UserModule(test.File(name="user_script.py",
data=user_script))
module.create_tmp_dir_with_files(tmp_code_dir)
hyperparameters = dict(
training_data_file=os.path.join(channel.path, "training_data.npz"),
sagemaker_program="user_script.py",
sagemaker_submit_directory=tmp_code_dir,
epochs=10,
batch_size=64,
model_dir=env.model_dir,
)
test.prepare(user_module=module,
hyperparameters=hyperparameters,
channels=[channel],
local=True)
assert execute_an_wrap_exit(
training_fn, capture_error=capture_error) == trainer.SUCCESS_CODE
model_path = os.path.join(env.model_dir, "saved_model")
model = fake_ml_framework.Model.load(model_path)
assert model.epochs == 10
assert model.batch_size == 64
assert model.loss == "elastic"
assert model.optimizer == "SGD"
def test_trainer_report_success(user_script):
with pytest.raises(ImportError):
importlib.import_module(modules.DEFAULT_MODULE_NAME)
channel = test.Channel.create(name='training')
features = [1, 2, 3, 4]
labels = [0, 1, 0, 1]
np.savez(os.path.join(channel.path, 'training_data'),
features=features,
labels=labels)
module = test.UserModule(test.File(name='user_script.py',
data=user_script))
hyperparameters = dict(training_data_file='training_data.npz',
sagemaker_program='user_script.py',
epochs=10,
batch_size=64,
optimizer='SGD')
test.prepare(user_module=module,
hyperparameters=hyperparameters,
channels=[channel])
os.environ[
'SAGEMAKER_TRAINING_MODULE'] = 'test.functional.simple_framework:train'
assert execute_an_wrap_exit(trainer.train) == trainer.SUCCESS_CODE
model_path = os.path.join(env.model_dir, 'saved_model')
model = fake_ml_framework.Model.load(model_path)
assert model.epochs == 10
assert model.batch_size == 64
assert model.optimizer == 'SGD'
assert os.path.exists(os.path.join(env.output_dir, 'success'))
def test_script_mode(user_script, training_fn, capture_error):
channel = test.Channel.create(name='training')
features = [1, 2, 3, 4]
labels = [0, 1, 0, 1]
np.savez(os.path.join(channel.path, 'training_data'), features=features, labels=labels)
module = test.UserModule(test.File(name='user_script.py', data=user_script))
hyperparameters = dict(training_data_file=os.path.join(channel.path, 'training_data.npz'),
sagemaker_program='user_script.py', epochs=10, batch_size=64, model_dir=env.model_dir)
test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
assert execute_an_wrap_exit(training_fn, capture_error=capture_error) == trainer.SUCCESS_CODE
model_path = os.path.join(env.model_dir, 'saved_model')
model = fake_ml_framework.Model.load(model_path)
assert model.epochs == 10
assert model.batch_size == 64
assert model.loss == 'elastic'
assert model.optimizer == 'SGD'
def test_trainer_report_failure():
channel = test.Channel.create(name="training")
features = [1, 2, 3, 4]
labels = [0, 1, 0, 1]
np.savez(os.path.join(channel.path, "training_data"),
features=features,
labels=labels)
module = test.UserModule(
test.File(name="user_script.py",
data=USER_SCRIPT_WITH_EXCEPTION)).add_file(setup_file)
hyperparameters = dict(
training_data_file="training_data.npz",
sagemaker_program="user_script.py",
epochs=10,
batch_size=64,
)
test.prepare(user_module=module,
hyperparameters=hyperparameters,
channels=[channel])
os.environ[
"SAGEMAKER_TRAINING_MODULE"] = "test.functional.simple_framework:train"
assert execute_an_wrap_exit(trainer.train) == errno.ENOENT
failure_file = os.path.join(env.output_dir, "failure")
assert os.path.exists(failure_file)
message = failure_message()
assert message.startswith("framework error:")
assert "No such file or directory" in message
def test_trainer_report_success(user_script, sagemaker_program):
channel = test.Channel.create(name='training')
features = [1, 2, 3, 4]
labels = [0, 1, 0, 1]
np.savez(os.path.join(channel.path, 'training_data'), features=features, labels=labels)
module = test.UserModule(test.File(name=sagemaker_program, data=user_script))
hyperparameters = dict(training_data_file='training_data.npz', sagemaker_program=sagemaker_program, epochs=10,
batch_size=64)
test.prepare(user_module=module, hyperparameters=hyperparameters, channels=[channel])
assert execute_an_wrap_exit(trainer.train) == trainer.SUCCESS_CODE
model_path = os.path.join(env.model_dir, 'saved_model')
model = fake_ml_framework.Model.load(model_path)
assert model.epochs == 10
assert model.batch_size == 64
assert model.optimizer == 'SGD'
assert os.path.exists(os.path.join(env.output_dir, 'success'))
import shlex
import subprocess
import textwrap
import pytest
import six
from sagemaker_containers.beta.framework import errors, modules
import test
data = [
"from distutils.core import setup\n",
'setup(name="my_test_script", py_modules=["my_test_script"])',
]
SETUP_FILE = test.File("setup.py", data)
USER_SCRIPT_FILE = test.File("my_test_script.py",
"def validate(): return True")
REQUIREMENTS_TXT_ASSERT_STR = """
____ __ __ _.............
/ ___| __ _ __ _ ___| \/ | __ _| | _____ _ __.
\___ \ / _` |/ _` |/ _ \ |\/| |/ _` | |/ / _ \ '__|
___) | (_| | (_| | __/ | | | (_| | < __/ |...
|____/ \__,_|\__, |\___|_| |_|\__,_|_|\_\___|_|...
|___/.................................
""".replace( # noqa W605
".", " ").strip()
import shlex
import subprocess
import textwrap
import pytest
from sagemaker_containers.beta.framework import errors, modules
import test
data = [
'from distutils.core import setup\n',
'setup(name="my_test_script", py_modules=["my_test_script"])'
]
SETUP = test.File('setup.py', data)
USER_SCRIPT = test.File('my_test_script.py', 'def validate(): return True')
@pytest.fixture(name='user_module_name')
def erase_user_module():
user_module = 'my_test_script'
yield user_module
try:
subprocess.check_call(
shlex.split('pip uninstall -y --quiet %s' % user_module))
except subprocess.CalledProcessError:
pass
import os
import shlex
import subprocess
import textwrap
import pytest
from sagemaker_containers.beta.framework import errors, modules
import test
data = [
"from distutils.core import setup\n",
'setup(name="my_test_script", py_modules=["my_test_script"])',
]
SETUP_FILE = test.File("setup.py", data)
USER_SCRIPT_FILE = test.File("my_test_script.py",
"def validate(): return True")
REQUIREMENTS_TXT_ASSERT_STR = """
____ __ __ _.............
/ ___| __ _ __ _ ___| \/ | __ _| | _____ _ __.
\___ \ / _` |/ _` |/ _ \ |\/| |/ _` | |/ / _ \ '__|
___) | (_| | (_| | __/ | | | (_| | < __/ |...
|____/ \__,_|\__, |\___|_| |_|\__,_|_|\_\___|_|...
|___/.................................
""".replace( # noqa W605
".", " ").strip()
import shlex
import subprocess
import textwrap
import pytest
import six
from sagemaker_containers.beta.framework import errors, modules
import test
data = [
'from distutils.core import setup\n',
'setup(name="my_test_script", py_modules=["my_test_script"])'
]
SETUP_FILE = test.File('setup.py', data)
USER_SCRIPT_FILE = test.File('my_test_script.py',
'def validate(): return True')
REQUIREMENTS_TXT_ASSERT_STR = """
____ __ __ _.............
/ ___| __ _ __ _ ___| \/ | __ _| | _____ _ __.
\___ \ / _` |/ _` |/ _ \ |\/| |/ _` | |/ / _ \ '__|
___) | (_| | (_| | __/ | | | (_| | < __/ |...
|____/ \__,_|\__, |\___|_| |_|\__,_|_|\_\___|_|...
|___/.................................
""".replace('.', ' ').strip() # noqa W605
@pytest.fixture(name='user_module_name')
BASH_SCRIPT = "#!/usr/bin/env python\n%s" % USER_MODE_SCRIPT
PARAMETER_SERVER_SCRIPT = """
from time import sleep
while True:
print('Running parameter server')
sleep(1)
"""
setup_file = test.File(
"setup.py",
"""
from setuptools import setup
setup(packages=[''],
name="user_script",
version='1.0.0',
include_package_data=True)
""",
)
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: