def test_mxnet_neo(strftime, sagemaker_session, mxnet_version):
mx = MXNet(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
framework_version=mxnet_version)
inputs = 's3://mybucket/train'
mx.fit(inputs=inputs)
input_shape = {'data': [100, 1, 28, 28]}
output_location = 's3://neo-sdk-test'
compiled_model = mx.compile_model(target_instance_family='ml_c4', input_shape=input_shape,
output_path=output_location)
sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
assert sagemaker_call_names == ['train', 'logs_for_job', 'sagemaker_client.describe_training_job',
'compile_model', 'wait_for_compilation_job']
expected_compile_model_args = _create_compilation_job(json.dumps(input_shape), output_location)
actual_compile_model_args = sagemaker_session.method_calls[3][2]
assert expected_compile_model_args == actual_compile_model_args
assert compiled_model.image == _neo_inference_image(mxnet_version)
predictor = mx.deploy(1, CPU, use_compiled_model=True)
assert isinstance(predictor, MXNetPredictor)
with pytest.raises(Exception) as wrong_target:
mx.deploy(1, CPU_C5, use_compiled_model=True)
assert str(wrong_target.value).startswith('No compiled model for')
# deploy without sagemaker Neo should continue to work
mx.deploy(1, CPU)
def test_mxnet_neo(strftime, sagemaker_session, neo_mxnet_version):
mx = MXNet(
entry_point=SCRIPT_PATH,
framework_version="1.6",
py_version="py3",
role=ROLE,
sagemaker_session=sagemaker_session,
instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
base_job_name="sagemaker-mxnet",
)
mx.fit()
input_shape = {"data": [100, 1, 28, 28]}
output_location = "s3://neo-sdk-test"
compiled_model = mx.compile_model(
target_instance_family="ml_c4",
input_shape=input_shape,
output_path=output_location,
framework="mxnet",
framework_version=neo_mxnet_version,
)
sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
assert sagemaker_call_names == [
"train",
"logs_for_job",
"sagemaker_client.describe_training_job",
"compile_model",
"wait_for_compilation_job",
]
expected_compile_model_args = _create_compilation_job(
json.dumps(input_shape), output_location)
actual_compile_model_args = sagemaker_session.method_calls[3][2]
assert expected_compile_model_args == actual_compile_model_args
assert compiled_model.image_uri == _neo_inference_image(neo_mxnet_version)
predictor = mx.deploy(1, CPU, use_compiled_model=True)
assert isinstance(predictor, MXNetPredictor)
with pytest.raises(Exception) as wrong_target:
mx.deploy(1, CPU_C5, use_compiled_model=True)
assert str(wrong_target.value).startswith("No compiled model for")
# deploy without sagemaker Neo should continue to work
mx.deploy(1, CPU)
