def test_create_model_with_optional_params(sagemaker_session,
coach_mxnet_version):
container_log_level = '"logging.INFO"'
source_dir = "s3://mybucket/source"
rl = RLEstimator(
entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
toolkit=RLToolkit.COACH,
toolkit_version=coach_mxnet_version,
framework=RLFramework.MXNET,
container_log_level=container_log_level,
source_dir=source_dir,
)
rl.fit(job_name="new_name")
new_role = "role"
new_entry_point = "deploy_script.py"
vpc_config = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]}
model_name = "model-name"
model = rl.create_model(role=new_role,
entry_point=new_entry_point,
vpc_config_override=vpc_config,
name=model_name)
assert model.role == new_role
assert model.vpc_config == vpc_config
assert model.entry_point == new_entry_point
assert model.name == model_name
def test_create_model_with_custom_image(name_from_base, sagemaker_session):
container_log_level = '"logging.INFO"'
source_dir = "s3://mybucket/source"
image = "selfdrivingcars:9000"
rl = RLEstimator(
entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
image_uri=image,
container_log_level=container_log_level,
source_dir=source_dir,
)
job_name = "new_name"
rl.fit(job_name=job_name)
model_name = "model_name"
name_from_base.return_value = model_name
new_entry_point = "deploy_script.py"
model = rl.create_model(entry_point=new_entry_point)
assert model.sagemaker_session == sagemaker_session
assert model.image_uri == image
assert model.entry_point == new_entry_point
assert model.role == ROLE
assert model.name == model_name
assert model.container_log_level == container_log_level
assert model.source_dir == source_dir
name_from_base.assert_called_with("selfdrivingcars")
def test_create_mxnet_model(sagemaker_session, rl_coach_mxnet_version):
container_log_level = '"logging.INFO"'
source_dir = "s3://mybucket/source"
rl = RLEstimator(
entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
train_instance_count=INSTANCE_COUNT,
train_instance_type=INSTANCE_TYPE,
toolkit=RLToolkit.COACH,
toolkit_version=rl_coach_mxnet_version,
framework=RLFramework.MXNET,
container_log_level=container_log_level,
source_dir=source_dir,
)
job_name = "new_name"
rl.fit(inputs="s3://mybucket/train", job_name="new_name")
model = rl.create_model()
supported_versions = TOOLKIT_FRAMEWORK_VERSION_MAP[RLToolkit.COACH.value]
framework_version = supported_versions[rl_coach_mxnet_version][
RLFramework.MXNET.value]
assert isinstance(model, MXNetModel)
assert model.sagemaker_session == sagemaker_session
assert model.framework_version == framework_version
assert model.py_version == PYTHON_VERSION
assert model.entry_point == SCRIPT_PATH
assert model.role == ROLE
assert model.name == job_name
assert model.container_log_level == container_log_level
assert model.source_dir == source_dir
assert model.vpc_config is None
def test_create_model_with_custom_image(sagemaker_session):
container_log_level = '"logging.INFO"'
source_dir = 's3://mybucket/source'
image = 'selfdrivingcars:9000'
rl = RLEstimator(entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
train_instance_count=INSTANCE_COUNT,
train_instance_type=INSTANCE_TYPE,
image_name=image,
container_log_level=container_log_level,
source_dir=source_dir)
job_name = 'new_name'
rl.fit(job_name=job_name)
new_entry_point = 'deploy_script.py'
model = rl.create_model(entry_point=new_entry_point)
assert model.sagemaker_session == sagemaker_session
assert model.image == image
assert model.entry_point == new_entry_point
assert model.role == ROLE
assert model.name == job_name
assert model.container_log_level == container_log_level
assert model.source_dir == source_dir
def test_create_model_with_optional_params(sagemaker_session,
rl_coach_mxnet_version):
container_log_level = '"logging.INFO"'
source_dir = 's3://mybucket/source'
rl = RLEstimator(entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
train_instance_count=INSTANCE_COUNT,
train_instance_type=INSTANCE_TYPE,
toolkit=RLToolkit.COACH,
toolkit_version=rl_coach_mxnet_version,
framework=RLFramework.MXNET,
container_log_level=container_log_level,
source_dir=source_dir)
rl.fit(job_name='new_name')
new_role = 'role'
new_entry_point = 'deploy_script.py'
vpc_config = {'Subnets': ['foo'], 'SecurityGroupIds': ['bar']}
model = rl.create_model(role=new_role,
entry_point=new_entry_point,
vpc_config_override=vpc_config)
assert model.role == new_role
assert model.vpc_config == vpc_config
assert model.entry_point == new_entry_point
def test_create_tf_model(sagemaker_session, rl_coach_tf_version):
container_log_level = '"logging.INFO"'
source_dir = 's3://mybucket/source'
rl = RLEstimator(entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
train_instance_count=INSTANCE_COUNT,
train_instance_type=INSTANCE_TYPE,
toolkit=RLToolkit.COACH,
toolkit_version=rl_coach_tf_version,
framework=RLFramework.TENSORFLOW,
container_log_level=container_log_level,
source_dir=source_dir)
job_name = 'new_name'
rl.fit(inputs='s3://mybucket/train', job_name='new_name')
model = rl.create_model()
supported_versions = TOOLKIT_FRAMEWORK_VERSION_MAP[RLToolkit.COACH.value]
framework_version = supported_versions[rl_coach_tf_version][
RLFramework.TENSORFLOW.value]
assert isinstance(model, tfs.Model)
assert model.sagemaker_session == sagemaker_session
assert model._framework_version == framework_version
assert model.role == ROLE
assert model.name == job_name
assert model._container_log_level == container_log_level
assert model.vpc_config is None
def test_create_tf_model(name_from_base, sagemaker_session, coach_tensorflow_version):
container_log_level = '"logging.INFO"'
source_dir = "s3://mybucket/source"
rl = RLEstimator(
entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
toolkit=RLToolkit.COACH,
toolkit_version=coach_tensorflow_version,
framework=RLFramework.TENSORFLOW,
container_log_level=container_log_level,
source_dir=source_dir,
)
rl.fit(inputs="s3://mybucket/train", job_name="new_name")
model_name = "model_name"
name_from_base.return_value = model_name
model = rl.create_model()
supported_versions = TOOLKIT_FRAMEWORK_VERSION_MAP[RLToolkit.COACH.value]
framework_version = supported_versions[coach_tensorflow_version][RLFramework.TENSORFLOW.value]
assert isinstance(model, TensorFlowModel)
assert model.sagemaker_session == sagemaker_session
assert model.framework_version == framework_version
assert model.role == ROLE
assert model.name == model_name
assert model._container_log_level == container_log_level
assert model.vpc_config is None
call_args = name_from_base.call_args_list[0][0]
assert call_args[0] in ("sagemaker-rl-tensorflow", "sagemaker-rl-coach-container")
def test_rl(strftime, sagemaker_session, rl_coach_mxnet_version):
rl = RLEstimator(
entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
train_instance_count=INSTANCE_COUNT,
train_instance_type=INSTANCE_TYPE,
toolkit=RLToolkit.COACH,
toolkit_version=rl_coach_mxnet_version,
framework=RLFramework.MXNET,
)
inputs = "s3://mybucket/train"
rl.fit(inputs=inputs, experiment_config=EXPERIMENT_CONFIG)
sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
assert sagemaker_call_names == ["train", "logs_for_job"]
boto_call_names = [
c[0] for c in sagemaker_session.boto_session.method_calls
]
assert boto_call_names == ["resource"]
expected_train_args = _create_train_job(RLToolkit.COACH.value,
rl_coach_mxnet_version,
RLFramework.MXNET.value)
expected_train_args["input_config"][0]["DataSource"]["S3DataSource"][
"S3Uri"] = inputs
expected_train_args["experiment_config"] = EXPERIMENT_CONFIG
actual_train_args = sagemaker_session.method_calls[0][2]
assert actual_train_args == expected_train_args
model = rl.create_model()
supported_versions = TOOLKIT_FRAMEWORK_VERSION_MAP[RLToolkit.COACH.value]
framework_version = supported_versions[rl_coach_mxnet_version][
RLFramework.MXNET.value]
expected_image_base = "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:{}-gpu-py3"
submit_dir = "s3://notmybucket/sagemaker-rl-mxnet-{}/source/sourcedir.tar.gz".format(
TIMESTAMP)
assert {
"Environment": {
"SAGEMAKER_SUBMIT_DIRECTORY": submit_dir,
"SAGEMAKER_PROGRAM": "dummy_script.py",
"SAGEMAKER_ENABLE_CLOUDWATCH_METRICS": "false",
"SAGEMAKER_REGION": "us-west-2",
"SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
},
"Image": expected_image_base.format(framework_version),
"ModelDataUrl": "s3://m/m.tar.gz",
} == model.prepare_container_def(GPU)
assert "cpu" in model.prepare_container_def(CPU)["Image"]
def test_rl(strftime, sagemaker_session, rl_coach_mxnet_version):
rl = RLEstimator(entry_point=SCRIPT_PATH,
role=ROLE,
sagemaker_session=sagemaker_session,
train_instance_count=INSTANCE_COUNT,
train_instance_type=INSTANCE_TYPE,
toolkit=RLToolkit.COACH,
toolkit_version=rl_coach_mxnet_version,
framework=RLFramework.MXNET)
inputs = 's3://mybucket/train'
rl.fit(inputs=inputs)
sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
assert sagemaker_call_names == ['train', 'logs_for_job']
boto_call_names = [
c[0] for c in sagemaker_session.boto_session.method_calls
]
assert boto_call_names == ['resource']
expected_train_args = _create_train_job(RLToolkit.COACH.value,
rl_coach_mxnet_version,
RLFramework.MXNET.value)
expected_train_args['input_config'][0]['DataSource']['S3DataSource'][
'S3Uri'] = inputs
actual_train_args = sagemaker_session.method_calls[0][2]
assert actual_train_args == expected_train_args
model = rl.create_model()
supported_versions = TOOLKIT_FRAMEWORK_VERSION_MAP[RLToolkit.COACH.value]
framework_version = supported_versions[rl_coach_mxnet_version][
RLFramework.MXNET.value]
expected_image_base = '520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:{}-gpu-py3'
submit_dir = 's3://notmybucket/sagemaker-rl-mxnet-{}/source/sourcedir.tar.gz'.format(
TIMESTAMP)
assert {
'Environment': {
'SAGEMAKER_SUBMIT_DIRECTORY': submit_dir,
'SAGEMAKER_PROGRAM': 'dummy_script.py',
'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
'SAGEMAKER_REGION': 'us-west-2',
'SAGEMAKER_CONTAINER_LOG_LEVEL': '20'
},
'Image': expected_image_base.format(framework_version),
'ModelDataUrl': 's3://m/m.tar.gz'
} == model.prepare_container_def(GPU)
assert 'cpu' in model.prepare_container_def(CPU)['Image']
