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']