def output_kms_key(sagemaker_session):
return get_or_create_kms_key(
sagemaker_session=sagemaker_session,
role_arn=ROLE,
alias="integ-test-processing-output-kms-key-{}".format(
sagemaker_session.boto_session.region_name),
)
def test_transform_mxnet(mxnet_estimator, mxnet_transform_input,
sagemaker_session, cpu_instance_type):
kms_key_arn = get_or_create_kms_key(sagemaker_session)
output_filter = "$"
input_filter = "$"
transformer = _create_transformer_and_transform_job(
mxnet_estimator,
mxnet_transform_input,
cpu_instance_type,
kms_key_arn,
input_filter=input_filter,
output_filter=output_filter,
join_source=None,
)
with timeout_and_delete_model_with_transformer(
transformer,
sagemaker_session,
minutes=TRANSFORM_DEFAULT_TIMEOUT_MINUTES):
transformer.wait()
job_desc = transformer.sagemaker_session.describe_transform_job(
job_name=transformer.latest_transform_job.name)
assert kms_key_arn == job_desc["TransformResources"]["VolumeKmsKeyId"]
assert output_filter == job_desc["DataProcessing"]["OutputFilter"]
assert input_filter == job_desc["DataProcessing"]["InputFilter"]
def test_transform_mxnet(sagemaker_session, mxnet_full_version):
data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
script_path = os.path.join(data_path, 'mnist.py')
mx = MXNet(entry_point=script_path, role='SageMakerRole', train_instance_count=1,
train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_session,
framework_version=mxnet_full_version)
train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
key_prefix='integ-test-data/mxnet_mnist/train')
test_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'test'),
key_prefix='integ-test-data/mxnet_mnist/test')
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
mx.fit({'train': train_input, 'test': test_input})
transform_input_path = os.path.join(data_path, 'transform', 'data.csv')
transform_input_key_prefix = 'integ-test-data/mxnet_mnist/transform'
transform_input = mx.sagemaker_session.upload_data(path=transform_input_path,
key_prefix=transform_input_key_prefix)
sts_client = sagemaker_session.boto_session.client('sts')
account_id = sts_client.get_caller_identity()['Account']
kms_client = sagemaker_session.boto_session.client('kms')
kms_key_arn = get_or_create_kms_key(kms_client, account_id)
transformer = _create_transformer_and_transform_job(mx, transform_input, kms_key_arn)
with timeout(minutes=TRANSFORM_DEFAULT_TIMEOUT_MINUTES):
transformer.wait()
job_desc = transformer.sagemaker_session.sagemaker_client.describe_transform_job(
TransformJobName=transformer.latest_transform_job.name)
assert kms_key_arn == job_desc['TransformResources']['VolumeKmsKeyId']
def volume_kms_key(sagemaker_session):
role_arn = sagemaker_session.expand_role(ROLE)
return get_or_create_kms_key(
sagemaker_session=sagemaker_session,
role_arn=role_arn,
alias="integ-test-processing-volume-kms-key-{}".format(
sagemaker_session.boto_session.region_name),
)
def test_deploy_model_with_tags_and_kms(
mxnet_training_job,
sagemaker_session,
mxnet_inference_latest_version,
mxnet_inference_latest_py_version,
cpu_instance_type,
):
endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
desc = sagemaker_session.sagemaker_client.describe_training_job(
TrainingJobName=mxnet_training_job
)
model_data = desc["ModelArtifacts"]["S3ModelArtifacts"]
script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist.py")
model = MXNetModel(
model_data,
"SageMakerRole",
entry_point=script_path,
py_version=mxnet_inference_latest_py_version,
sagemaker_session=sagemaker_session,
framework_version=mxnet_inference_latest_version,
)
tags = [{"Key": "TagtestKey", "Value": "TagtestValue"}]
kms_key_arn = get_or_create_kms_key(sagemaker_session)
model.deploy(
1, cpu_instance_type, endpoint_name=endpoint_name, tags=tags, kms_key=kms_key_arn
)
returned_model = sagemaker_session.sagemaker_client.describe_model(ModelName=model.name)
returned_model_tags = sagemaker_session.sagemaker_client.list_tags(
ResourceArn=returned_model["ModelArn"]
)["Tags"]
endpoint = sagemaker_session.sagemaker_client.describe_endpoint(EndpointName=endpoint_name)
endpoint_tags = sagemaker_session.sagemaker_client.list_tags(
ResourceArn=endpoint["EndpointArn"]
)["Tags"]
endpoint_config = sagemaker_session.sagemaker_client.describe_endpoint_config(
EndpointConfigName=endpoint["EndpointConfigName"]
)
endpoint_config_tags = sagemaker_session.sagemaker_client.list_tags(
ResourceArn=endpoint_config["EndpointConfigArn"]
)["Tags"]
production_variants = endpoint_config["ProductionVariants"]
assert returned_model_tags == tags
assert endpoint_config_tags == tags
assert endpoint_tags == tags
assert production_variants[0]["InstanceType"] == cpu_instance_type
assert production_variants[0]["InitialInstanceCount"] == 1
assert endpoint_config["KmsKeyId"] == kms_key_arn
def test_transform_mxnet(sagemaker_session, mxnet_full_version, cpu_instance_type):
data_path = os.path.join(DATA_DIR, "mxnet_mnist")
script_path = os.path.join(data_path, "mnist.py")
mx = MXNet(
entry_point=script_path,
role="SageMakerRole",
train_instance_count=1,
train_instance_type=cpu_instance_type,
sagemaker_session=sagemaker_session,
framework_version=mxnet_full_version,
)
train_input = mx.sagemaker_session.upload_data(
path=os.path.join(data_path, "train"), key_prefix="integ-test-data/mxnet_mnist/train"
)
test_input = mx.sagemaker_session.upload_data(
path=os.path.join(data_path, "test"), key_prefix="integ-test-data/mxnet_mnist/test"
)
job_name = unique_name_from_base("test-mxnet-transform")
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
mx.fit({"train": train_input, "test": test_input}, job_name=job_name)
transform_input_path = os.path.join(data_path, "transform", "data.csv")
transform_input_key_prefix = "integ-test-data/mxnet_mnist/transform"
transform_input = mx.sagemaker_session.upload_data(
path=transform_input_path, key_prefix=transform_input_key_prefix
)
kms_key_arn = get_or_create_kms_key(sagemaker_session)
output_filter = "$"
input_filter = "$"
transformer = _create_transformer_and_transform_job(
mx,
transform_input,
cpu_instance_type,
kms_key_arn,
input_filter=input_filter,
output_filter=output_filter,
join_source=None,
)
with timeout_and_delete_model_with_transformer(
transformer, sagemaker_session, minutes=TRANSFORM_DEFAULT_TIMEOUT_MINUTES
):
transformer.wait()
job_desc = transformer.sagemaker_session.sagemaker_client.describe_transform_job(
TransformJobName=transformer.latest_transform_job.name
)
assert kms_key_arn == job_desc["TransformResources"]["VolumeKmsKeyId"]
assert output_filter == job_desc["DataProcessing"]["OutputFilter"]
assert input_filter == job_desc["DataProcessing"]["InputFilter"]
def test_deploy_model_with_tags_and_kms(mxnet_training_job, sagemaker_session,
mxnet_full_version):
endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
desc = sagemaker_session.sagemaker_client.describe_training_job(
TrainingJobName=mxnet_training_job)
model_data = desc['ModelArtifacts']['S3ModelArtifacts']
script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
model = MXNetModel(model_data,
'SageMakerRole',
entry_point=script_path,
py_version=PYTHON_VERSION,
sagemaker_session=sagemaker_session,
framework_version=mxnet_full_version)
tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
kms_key_arn = get_or_create_kms_key(sagemaker_session)
model.deploy(1,
'ml.m4.xlarge',
endpoint_name=endpoint_name,
tags=tags,
kms_key=kms_key_arn)
returned_model = sagemaker_session.describe_model(
EndpointName=model.name)
returned_model_tags = sagemaker_session.list_tags(
ResourceArn=returned_model['ModelArn'])['Tags']
endpoint = sagemaker_session.describe_endpoint(
EndpointName=endpoint_name)
endpoint_tags = sagemaker_session.list_tags(
ResourceArn=endpoint['EndpointArn'])['Tags']
endpoint_config = sagemaker_session.describe_endpoint_config(
EndpointConfigName=endpoint['EndpointConfigName'])
endpoint_config_tags = sagemaker_session.list_tags(
ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
production_variants = endpoint_config['ProductionVariants']
assert returned_model_tags == tags
assert endpoint_config_tags == tags
assert endpoint_tags == tags
assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
assert production_variants[0]['InitialInstanceCount'] == 1
assert endpoint_config['KmsKeyId'] == kms_key_arn
def test_transform_mxnet(sagemaker_session, mxnet_full_version):
data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
script_path = os.path.join(data_path, 'mnist.py')
mx = MXNet(entry_point=script_path,
role='SageMakerRole',
train_instance_count=1,
train_instance_type='ml.c4.xlarge',
sagemaker_session=sagemaker_session,
framework_version=mxnet_full_version)
train_input = mx.sagemaker_session.upload_data(
path=os.path.join(data_path, 'train'),
key_prefix='integ-test-data/mxnet_mnist/train')
test_input = mx.sagemaker_session.upload_data(
path=os.path.join(data_path, 'test'),
key_prefix='integ-test-data/mxnet_mnist/test')
job_name = unique_name_from_base('test-mxnet-transform')
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
mx.fit({'train': train_input, 'test': test_input}, job_name=job_name)
transform_input_path = os.path.join(data_path, 'transform', 'data.csv')
transform_input_key_prefix = 'integ-test-data/mxnet_mnist/transform'
transform_input = mx.sagemaker_session.upload_data(
path=transform_input_path, key_prefix=transform_input_key_prefix)
kms_key_arn = get_or_create_kms_key(sagemaker_session)
output_filter = "$"
transformer = _create_transformer_and_transform_job(
mx,
transform_input,
kms_key_arn,
input_filter=None,
output_filter=output_filter,
join_source=None)
with timeout_and_delete_model_with_transformer(
transformer,
sagemaker_session,
minutes=TRANSFORM_DEFAULT_TIMEOUT_MINUTES):
transformer.wait()
job_desc = transformer.sagemaker_session.sagemaker_client.describe_transform_job(
TransformJobName=transformer.latest_transform_job.name)
assert kms_key_arn == job_desc['TransformResources']['VolumeKmsKeyId']
assert output_filter == job_desc['DataProcessing']['OutputFilter']
def monitoring_files_kms_key(sagemaker_session):
return get_or_create_kms_key(sagemaker_session=sagemaker_session)
def s3_files_kms_key(sagemaker_session):
return get_or_create_kms_key(sagemaker_session=sagemaker_session)
def test_model_registration_with_model_repack(
sagemaker_session,
role,
pipeline_name,
region_name,
):
kms_key = get_or_create_kms_key(sagemaker_session, role)
base_dir = os.path.join(DATA_DIR, "pytorch_mnist")
entry_point = os.path.join(base_dir, "mnist.py")
input_path = sagemaker_session.upload_data(
path=os.path.join(base_dir, "training"),
key_prefix="integ-test-data/pytorch_mnist/training",
)
inputs = TrainingInput(s3_data=input_path)
instance_count = ParameterInteger(name="InstanceCount", default_value=1)
instance_type = ParameterString(name="InstanceType",
default_value="ml.m5.xlarge")
good_enough_input = ParameterInteger(name="GoodEnoughInput",
default_value=1)
pytorch_estimator = PyTorch(
entry_point=entry_point,
role=role,
framework_version="1.5.0",
py_version="py3",
instance_count=instance_count,
instance_type=instance_type,
sagemaker_session=sagemaker_session,
output_kms_key=kms_key,
)
step_train = TrainingStep(
name="pytorch-train",
estimator=pytorch_estimator,
inputs=inputs,
)
step_register = RegisterModel(
name="pytorch-register-model",
estimator=pytorch_estimator,
model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
content_types=["text/csv"],
response_types=["text/csv"],
inference_instances=["ml.t2.medium", "ml.m5.large"],
transform_instances=["ml.m5.large"],
description="test-description",
entry_point=entry_point,
model_kms_key=kms_key,
)
model = Model(
image_uri=pytorch_estimator.training_image_uri(),
model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
sagemaker_session=sagemaker_session,
role=role,
)
model_inputs = CreateModelInput(
instance_type="ml.m5.large",
accelerator_type="ml.eia1.medium",
)
step_model = CreateModelStep(
name="pytorch-model",
model=model,
inputs=model_inputs,
)
step_cond = ConditionStep(
name="cond-good-enough",
conditions=[
ConditionGreaterThanOrEqualTo(left=good_enough_input, right=1)
],
if_steps=[step_train, step_register],
else_steps=[step_model],
)
pipeline = Pipeline(
name=pipeline_name,
parameters=[good_enough_input, instance_count, instance_type],
steps=[step_cond],
sagemaker_session=sagemaker_session,
)
try:
response = pipeline.create(role)
create_arn = response["PipelineArn"]
assert re.match(
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
create_arn,
)
execution = pipeline.start(parameters={})
assert re.match(
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
execution.arn,
)
execution = pipeline.start(parameters={"GoodEnoughInput": 0})
assert re.match(
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
execution.arn,
)
finally:
try:
pipeline.delete()
except Exception:
pass
