def network_config():
return NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
encrypt_inter_container_traffic=True,
)
def script_processor(sagemaker_session):
return ScriptProcessor(
role=ROLE,
image_uri=
"012345678901.dkr.ecr.us-west-2.amazonaws.com/my-custom-image-uri",
command=["python3"],
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="my_sklearn_processor",
env={"my_env_variable": "my_env_variable_value"},
tags=[{
"Key": "my-tag",
"Value": "my-tag-value"
}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
encrypt_inter_container_traffic=True,
),
sagemaker_session=sagemaker_session,
)
def test_sklearn_with_network_config(sagemaker_session, sklearn_full_version,
cpu_instance_type):
script_path = os.path.join(DATA_DIR, "dummy_script.py")
input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
sklearn_processor = SKLearnProcessor(
framework_version=sklearn_full_version,
role=ROLE,
instance_type=cpu_instance_type,
instance_count=1,
command=["python3"],
sagemaker_session=sagemaker_session,
base_job_name="test-sklearn-with-network-config",
network_config=NetworkConfig(enable_network_isolation=True,
encrypt_inter_container_traffic=True),
)
sklearn_processor.run(
code=script_path,
inputs=[
ProcessingInput(source=input_file_path,
destination="/opt/ml/processing/inputs/")
],
wait=False,
logs=False,
)
job_description = sklearn_processor.latest_job.describe()
network_config = job_description["NetworkConfig"]
assert network_config["EnableInterContainerTrafficEncryption"]
assert network_config["EnableNetworkIsolation"]
def test_processor_with_all_parameters(sagemaker_session):
processor = Processor(
role=ROLE,
image_uri=CUSTOM_IMAGE_URI,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
entrypoint=[
"python3", "/opt/ml/processing/input/code/processing_code.py"
],
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="processor_base_name",
env={"my_env_variable": "my_env_variable_value"},
tags=[{
"Key": "my-tag",
"Value": "my-tag-value"
}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
),
)
processor.run(
inputs=[
ProcessingInput(
source="s3://path/to/my/dataset/census.csv",
destination="/container/path/",
input_name="my_dataset",
s3_data_type="S3Prefix",
s3_input_mode="File",
s3_data_distribution_type="FullyReplicated",
s3_compression_type="None",
)
],
outputs=[
ProcessingOutput(
source="/container/path/",
destination="s3://uri/",
output_name="my_output",
s3_upload_mode="EndOfJob",
)
],
arguments=["--drop-columns", "'SelfEmployed'"],
wait=True,
logs=False,
job_name="my_job_name",
experiment_config={"ExperimentName": "AnExperiment"},
)
expected_args = _get_expected_args_all_parameters(
processor._current_job_name)
# Drop the "code" input from expected values.
expected_args["inputs"] = [expected_args["inputs"][0]]
sagemaker_session.process.assert_called_with(**expected_args)
def test_sklearn_with_all_parameters_via_run_args_called_twice(
exists_mock, isfile_mock, botocore_resolver, sklearn_version, sagemaker_session
):
botocore_resolver.return_value.construct_endpoint.return_value = {"hostname": ECR_HOSTNAME}
processor = SKLearnProcessor(
role=ROLE,
framework_version=sklearn_version,
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="my_sklearn_processor",
env={"my_env_variable": "my_env_variable_value"},
tags=[{"Key": "my-tag", "Value": "my-tag-value"}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
encrypt_inter_container_traffic=True,
),
sagemaker_session=sagemaker_session,
)
run_args = processor.get_run_args(
code="/local/path/to/processing_code.py",
inputs=_get_data_inputs_all_parameters(),
outputs=_get_data_outputs_all_parameters(),
arguments=["--drop-columns", "'SelfEmployed'"],
)
run_args = processor.get_run_args(
code="/local/path/to/processing_code.py",
inputs=_get_data_inputs_all_parameters(),
outputs=_get_data_outputs_all_parameters(),
arguments=["--drop-columns", "'SelfEmployed'"],
)
processor.run(
code=run_args.code,
inputs=run_args.inputs,
outputs=run_args.outputs,
arguments=run_args.arguments,
wait=True,
logs=False,
experiment_config={"ExperimentName": "AnExperiment"},
)
expected_args = _get_expected_args_all_parameters(processor._current_job_name)
sklearn_image_uri = (
"246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:{}-cpu-py3"
).format(sklearn_version)
expected_args["app_specification"]["ImageUri"] = sklearn_image_uri
sagemaker_session.process.assert_called_with(**expected_args)
def test_processing_step_normalizes_args(mock_normalize_args,
sagemaker_session):
processor = ScriptProcessor(
role=ROLE,
image_uri=
"012345678901.dkr.ecr.us-west-2.amazonaws.com/my-custom-image-uri",
command=["python3"],
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="my_sklearn_processor",
env={"my_env_variable": "my_env_variable_value"},
tags=[{
"Key": "my-tag",
"Value": "my-tag-value"
}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
encrypt_inter_container_traffic=True,
),
sagemaker_session=sagemaker_session,
)
cache_config = CacheConfig(enable_caching=True, expire_after="PT1H")
inputs = [
ProcessingInput(
source=f"s3://{BUCKET}/processing_manifest",
destination="processing_manifest",
)
]
outputs = [
ProcessingOutput(
source=f"s3://{BUCKET}/processing_manifest",
destination="processing_manifest",
)
]
step = ProcessingStep(
name="MyProcessingStep",
processor=processor,
code="foo.py",
inputs=inputs,
outputs=outputs,
job_arguments=["arg1", "arg2"],
cache_config=cache_config,
)
mock_normalize_args.return_value = [step.inputs, step.outputs]
step.to_request()
mock_normalize_args.assert_called_with(
arguments=step.job_arguments,
inputs=step.inputs,
outputs=step.outputs,
code=step.code,
)
def test_script_processor_with_all_parameters(exists_mock, isfile_mock, sagemaker_session):
processor = ScriptProcessor(
role=ROLE,
image_uri=CUSTOM_IMAGE_URI,
command=["python3"],
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="my_sklearn_processor",
env={"my_env_variable": "my_env_variable_value"},
tags=[{"Key": "my-tag", "Value": "my-tag-value"}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
encrypt_inter_container_traffic=True,
),
sagemaker_session=sagemaker_session,
)
processor.run(
code="/local/path/to/processing_code.py",
inputs=[
ProcessingInput(
source="s3://path/to/my/dataset/census.csv",
destination="/container/path/",
input_name="my_dataset",
s3_data_type="S3Prefix",
s3_input_mode="File",
s3_data_distribution_type="FullyReplicated",
s3_compression_type="None",
)
],
outputs=[
ProcessingOutput(
app_managed=True,
feature_store_output=FeatureStoreOutput("Foo"),
source="/container/path/",
destination="s3://uri/",
output_name="my_output",
s3_upload_mode="EndOfJob",
)
],
arguments=["--drop-columns", "'SelfEmployed'"],
wait=True,
logs=False,
job_name="my_job_name",
experiment_config={"ExperimentName": "AnExperiment"},
)
expected_args = _get_expected_args_all_parameters(processor._current_job_name)
sagemaker_session.process.assert_called_with(**expected_args)
assert "my_job_name" in processor._current_job_name
def test_sklearn_with_all_parameters(exists_mock, isfile_mock, sagemaker_session):
processor = SKLearnProcessor(
role=ROLE,
framework_version="0.20.0",
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="my_sklearn_processor",
env={"my_env_variable": "my_env_variable_value"},
tags=[{"Key": "my-tag", "Value": "my-tag-value"}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
),
sagemaker_session=sagemaker_session,
)
processor.run(
code="/local/path/to/processing_code.py",
inputs=[
ProcessingInput(
source="s3://path/to/my/dataset/census.csv",
destination="/container/path/",
input_name="my_dataset",
s3_data_type="S3Prefix",
s3_input_mode="File",
s3_data_distribution_type="FullyReplicated",
s3_compression_type="None",
)
],
outputs=[
ProcessingOutput(
source="/container/path/",
destination="s3://uri/",
output_name="my_output",
s3_upload_mode="EndOfJob",
)
],
arguments=["--drop-columns", "'SelfEmployed'"],
wait=True,
logs=False,
job_name="my_job_name",
experiment_config={"ExperimentName": "AnExperiment"},
)
expected_args = _get_expected_args_all_parameters(processor._current_job_name)
sklearn_image_uri = (
"246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3"
)
expected_args["app_specification"]["ImageUri"] = sklearn_image_uri
sagemaker_session.process.assert_called_with(**expected_args)
def test_default_model_monitor_with_invalid_network_config(sagemaker_session):
invalid_network_config = NetworkConfig(
encrypt_inter_container_traffic=False)
my_default_monitor = DefaultModelMonitor(
role=ROLE,
sagemaker_session=sagemaker_session,
network_config=invalid_network_config)
with pytest.raises(ValueError) as exception:
my_default_monitor.create_monitoring_schedule(
endpoint_input="test_endpoint")
assert INTER_CONTAINER_ENCRYPTION_EXCEPTION_MSG in str(exception.value)
with pytest.raises(ValueError) as exception:
my_default_monitor.update_monitoring_schedule()
assert INTER_CONTAINER_ENCRYPTION_EXCEPTION_MSG in str(exception.value)
def processing(self) -> Kwargs:
"""Return the mandatory kwargs for `Processing`."""
if not hasattr(self, "_processing"):
self._processing = dict(
role=self.role,
tags=self.tags,
volume_kms_key=self.volume_kms_key,
output_kms_key=self.s3_kms_key,
network_config=NetworkConfig(
enable_network_isolation=self.network_isolation,
security_group_ids=self.sgs,
subnets=self.subnets,
),
)
return self._processing
def test_script_processor_with_all_parameters_via_run_args(
exists_mock, isfile_mock, sagemaker_session
):
processor = ScriptProcessor(
role=ROLE,
image_uri=CUSTOM_IMAGE_URI,
command=["python3"],
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="my_sklearn_processor",
env={"my_env_variable": "my_env_variable_value"},
tags=[{"Key": "my-tag", "Value": "my-tag-value"}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
encrypt_inter_container_traffic=True,
),
sagemaker_session=sagemaker_session,
)
run_args = processor.get_run_args(
code="/local/path/to/processing_code.py",
inputs=_get_data_inputs_all_parameters(),
outputs=_get_data_outputs_all_parameters(),
arguments=["--drop-columns", "'SelfEmployed'"],
)
processor.run(
code=run_args.code,
inputs=run_args.inputs,
outputs=run_args.outputs,
arguments=run_args.arguments,
wait=True,
logs=False,
job_name="my_job_name",
experiment_config={"ExperimentName": "AnExperiment"},
)
expected_args = _get_expected_args_all_parameters(processor._current_job_name)
sagemaker_session.process.assert_called_with(**expected_args)
assert "my_job_name" in processor._current_job_name
def test_extend_processing_args(sagemaker_session):
inputs = []
outputs = []
processor = Processor(
role=ROLE,
image_uri=CUSTOM_IMAGE_URI,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
network_config=NetworkConfig(encrypt_inter_container_traffic=False),
)
extended_inputs, extended_outputs = processor._extend_processing_args([], [])
assert extended_inputs == inputs
assert extended_outputs == outputs
def test_processor_with_missing_network_config_parameters(sagemaker_session):
processor = Processor(
role=ROLE,
image_uri=CUSTOM_IMAGE_URI,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
network_config=NetworkConfig(enable_network_isolation=True),
)
processor.run()
expected_args = _get_expected_args(processor._current_job_name)
del expected_args["app_specification"]["ContainerEntrypoint"]
expected_args["inputs"] = []
expected_args["network_config"] = {"EnableNetworkIsolation": True}
sagemaker_session.process.assert_called_with(**expected_args)
def test_processor_with_all_parameters(sagemaker_session):
processor = Processor(
role=ROLE,
image_uri=CUSTOM_IMAGE_URI,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
entrypoint=[
"python3", "/opt/ml/processing/input/code/processing_code.py"
],
volume_size_in_gb=100,
volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
max_runtime_in_seconds=3600,
base_job_name="processor_base_name",
env={"my_env_variable": "my_env_variable_value"},
tags=[{
"Key": "my-tag",
"Value": "my-tag-value"
}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
encrypt_inter_container_traffic=True,
),
)
processor.run(
inputs=_get_data_inputs_all_parameters(),
outputs=_get_data_outputs_all_parameters(),
arguments=["--drop-columns", "'SelfEmployed'"],
wait=True,
logs=False,
job_name="my_job_name",
experiment_config={"ExperimentName": "AnExperiment"},
)
expected_args = _get_expected_args_all_parameters(
processor._current_job_name)
# Drop the "code" input from expected values.
expected_args["inputs"] = expected_args["inputs"][:-1]
sagemaker_session.process.assert_called_with(**expected_args)
def test_model_monitor_with_invalid_network_config(sagemaker_session):
invalid_network_config = NetworkConfig(
encrypt_inter_container_traffic=False)
my_model_monitor = ModelMonitor(
role=ROLE,
image_uri=CUSTOM_IMAGE_URI,
sagemaker_session=sagemaker_session,
network_config=invalid_network_config,
)
with pytest.raises(ValueError) as exception:
my_model_monitor.create_monitoring_schedule(
endpoint_input="test_endpoint",
output=MonitoringOutput(source="/opt/ml/processing/output",
destination="/opt/ml/processing/output"),
)
assert INTER_CONTAINER_ENCRYPTION_EXCEPTION_MSG in str(exception.value)
with pytest.raises(ValueError) as exception:
my_model_monitor.update_monitoring_schedule()
assert INTER_CONTAINER_ENCRYPTION_EXCEPTION_MSG in str(exception.value)
def test_sklearn_with_all_customizations(sagemaker_session):
sklearn_processor = SKLearnProcessor(
framework_version="0.20.0",
role=ROLE,
command=["python3"],
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=100,
volume_kms_key=None,
output_kms_key="arn:aws:kms:us-west-2:012345678901:key/kms-key",
max_runtime_in_seconds=3600,
base_job_name="my_sklearn_processor",
env={"my_env_variable": "my_env_variable_value"},
tags=[{
"Key": "my-tag",
"Value": "my-tag-value"
}],
network_config=NetworkConfig(
subnets=["my_subnet_id"],
security_group_ids=["my_security_group_id"],
enable_network_isolation=True,
),
sagemaker_session=sagemaker_session,
)
with patch("os.path.isdir", return_value=True):
sklearn_processor.run(
code="/local/path/to/sklearn_transformer.py",
inputs=[
ProcessingInput(
source="s3://path/to/my/dataset/census.csv",
destination="/container/path/",
input_name="my_dataset",
s3_data_type="S3Prefix",
s3_input_mode="File",
s3_data_distribution_type="FullyReplicated",
s3_compression_type="None",
)
],
outputs=[
ProcessingOutput(
source="/container/path/",
destination="s3://uri/",
output_name="my_output",
s3_upload_mode="EndOfJob",
)
],
arguments=["--drop-columns", "'SelfEmployed'"],
wait=True,
logs=False,
job_name="my_job_name",
experiment_config={"ExperimentName": "AnExperiment"},
)
expected_args = {
"inputs": [
{
"InputName": "my_dataset",
"S3Input": {
"S3Uri": "s3://path/to/my/dataset/census.csv",
"LocalPath": "/container/path/",
"S3DataType": "S3Prefix",
"S3InputMode": "File",
"S3DataDistributionType": "FullyReplicated",
"S3CompressionType": "None",
},
},
{
"InputName": "code",
"S3Input": {
"S3Uri": "mocked_s3_uri_from_upload_data",
"LocalPath": "/opt/ml/processing/input/code",
"S3DataType": "S3Prefix",
"S3InputMode": "File",
"S3DataDistributionType": "FullyReplicated",
"S3CompressionType": "None",
},
},
],
"output_config": {
"Outputs": [{
"OutputName": "my_output",
"S3Output": {
"S3Uri": "s3://uri/",
"LocalPath": "/container/path/",
"S3UploadMode": "EndOfJob",
},
}],
"KmsKeyId":
"arn:aws:kms:us-west-2:012345678901:key/kms-key",
},
"job_name":
sklearn_processor._current_job_name,
"resources": {
"ClusterConfig": {
"InstanceType": "ml.m4.xlarge",
"InstanceCount": 1,
"VolumeSizeInGB": 100,
}
},
"stopping_condition": {
"MaxRuntimeInSeconds": 3600
},
"app_specification": {
"ImageUri":
"246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3",
"ContainerArguments": ["--drop-columns", "'SelfEmployed'"],
"ContainerEntrypoint": [
"python3",
"/opt/ml/processing/input/code/sklearn_transformer.py",
],
},
"environment": {
"my_env_variable": "my_env_variable_value"
},
"network_config": {
"EnableNetworkIsolation": True,
"VpcConfig": {
"SecurityGroupIds": ["my_security_group_id"],
"Subnets": ["my_subnet_id"],
},
},
"role_arn":
ROLE,
"tags": [{
"Key": "my-tag",
"Value": "my-tag-value"
}],
"experiment_config": {
"ExperimentName": "AnExperiment"
},
}
sagemaker_session.process.assert_called_with(**expected_args)
def get_pipeline(
region,
sagemaker_project_arn=None,
role=None,
default_bucket=None,
model_package_group_name="restatePackageGroup", # Choose any name
pipeline_name="restate-p-XXXXXXXXX", # You can find your pipeline name in the Studio UI (project -> Pipelines -> name)
base_job_prefix="restate", # Choose any name
):
"""Gets a SageMaker ML Pipeline instance working with on RE data.
Args:
region: AWS region to create and run the pipeline.
role: IAM role to create and run steps and pipeline.
default_bucket: the bucket to use for storing the artifacts
Returns:
an instance of a pipeline
"""
sagemaker_session = get_session(region, default_bucket)
if role is None:
role = sagemaker.session.get_execution_role(sagemaker_session)
# Parameters for pipeline execution
processing_instance_count = ParameterInteger(name="ProcessingInstanceCount", default_value=1)
processing_instance_type = ParameterString(
name="ProcessingInstanceType", default_value="ml.m5.2xlarge"
)
training_instance_type = ParameterString(
name="TrainingInstanceType", default_value="ml.m5.xlarge"
)
model_approval_status = ParameterString(
name="ModelApprovalStatus",
default_value="PendingManualApproval", # ModelApprovalStatus can be set to a default of "Approved" if you don't want manual approval.
)
input_data = ParameterString(
name="InputDataUrl",
default_value=f"", # Change this to point to the s3 location of your raw input data.
)
data_sources = []
# Sagemaker session
sess = sagemaker_session
# You can configure this with your own bucket name, e.g.
# bucket = "my-bucket"
bucket = sess.default_bucket()
data_sources.append(
ProcessingInput(
input_name="restate-california",
dataset_definition=DatasetDefinition(
local_path="/opt/ml/processing/restate-california",
data_distribution_type="FullyReplicated",
# You can override below to point to other database or use different queries
athena_dataset_definition=AthenaDatasetDefinition(
catalog="AwsDataCatalog",
database="restate",
query_string="SELECT * FROM restate.california_10",
output_s3_uri=f"s3://{bucket}/athena/",
output_format="PARQUET",
),
),
)
)
print(f"Data Wrangler export storage bucket: {bucket}")
# unique flow export ID
flow_export_id = f"{time.strftime('%d-%H-%M-%S', time.gmtime())}-{str(uuid.uuid4())[:8]}"
flow_export_name = f"flow-{flow_export_id}"
# Output name is auto-generated from the select node's ID + output name from the flow file.
output_name = "99ae1ec3-dd5f-453c-bfae-721dac423cd7.default"
s3_output_prefix = f"export-{flow_export_name}/output"
s3_output_path = f"s3://{bucket}/{s3_output_prefix}"
print(f"Flow S3 export result path: {s3_output_path}")
processing_job_output = ProcessingOutput(
output_name=output_name,
source="/opt/ml/processing/output",
destination=s3_output_path,
s3_upload_mode="EndOfJob",
)
# name of the flow file which should exist in the current notebook working directory
flow_file_name = "sagemaker-pipeline/restate-athena-california.flow"
# Load .flow file from current notebook working directory
#!echo "Loading flow file from current notebook working directory: $PWD"
with open(flow_file_name) as f:
flow = json.load(f)
# Upload flow to S3
s3_client = boto3.client("s3")
s3_client.upload_file(
flow_file_name,
bucket,
f"data_wrangler_flows/{flow_export_name}.flow",
ExtraArgs={"ServerSideEncryption": "aws:kms"},
)
flow_s3_uri = f"s3://{bucket}/data_wrangler_flows/{flow_export_name}.flow"
print(f"Data Wrangler flow {flow_file_name} uploaded to {flow_s3_uri}")
## Input - Flow: restate-athena-russia.flow
flow_input = ProcessingInput(
source=flow_s3_uri,
destination="/opt/ml/processing/flow",
input_name="flow",
s3_data_type="S3Prefix",
s3_input_mode="File",
s3_data_distribution_type="FullyReplicated",
)
# IAM role for executing the processing job.
iam_role = role
# Unique processing job name. Give a unique name every time you re-execute processing jobs
processing_job_name = f"data-wrangler-flow-processing-{flow_export_id}"
# Data Wrangler Container URL.
container_uri = sagemaker.image_uris.retrieve(
framework="data-wrangler", # we are using the Sagemaker built in xgboost algorithm
region=region,
)
# Processing Job Instance count and instance type.
instance_count = 2
instance_type = "ml.m5.4xlarge"
# Size in GB of the EBS volume to use for storing data during processing
volume_size_in_gb = 30
# Content type for each output. Data Wrangler supports CSV as default and Parquet.
output_content_type = "CSV"
# Network Isolation mode; default is off
enable_network_isolation = False
# List of tags to be passed to the processing job
user_tags = []
# Output configuration used as processing job container arguments
output_config = {output_name: {"content_type": output_content_type}}
# KMS key for per object encryption; default is None
kms_key = None
processor = Processor(
role=iam_role,
image_uri=container_uri,
instance_count=instance_count,
instance_type=instance_type,
volume_size_in_gb=volume_size_in_gb,
network_config=NetworkConfig(enable_network_isolation=enable_network_isolation),
sagemaker_session=sess,
output_kms_key=kms_key,
tags=user_tags,
)
data_wrangler_step = ProcessingStep(
name="DataWranglerProcess",
processor=processor,
inputs=[flow_input] + data_sources,
outputs=[processing_job_output],
job_arguments=[f"--output-config '{json.dumps(output_config)}'"],
)
# Processing step for feature engineering
# this processor does not have awswrangler installed
sklearn_processor = SKLearnProcessor(
framework_version="0.23-1",
instance_type=processing_instance_type,
instance_count=processing_instance_count,
base_job_name=f"{base_job_prefix}/sklearn-restate-preprocess", # choose any name
sagemaker_session=sagemaker_session,
role=role,
)
step_process = ProcessingStep(
name="Preprocess", # choose any name
processor=sklearn_processor,
inputs=[
ProcessingInput(
source=data_wrangler_step.properties.ProcessingOutputConfig.Outputs[
output_name
].S3Output.S3Uri,
destination="/opt/ml/processing/data/raw-data-dir",
)
],
outputs=[
ProcessingOutput(output_name="train", source="/opt/ml/processing/train"),
ProcessingOutput(output_name="validation", source="/opt/ml/processing/validation"),
ProcessingOutput(output_name="test", source="/opt/ml/processing/test"),
],
code=os.path.join(BASE_DIR, "preprocess.py"),
job_arguments=[
"--input-data",
data_wrangler_step.properties.ProcessingOutputConfig.Outputs[
output_name
].S3Output.S3Uri,
],
)
# Training step for generating model artifacts
model_path = f"s3://{sagemaker_session.default_bucket()}/{base_job_prefix}/restateTrain"
model_bucket_key = f"{sagemaker_session.default_bucket()}/{base_job_prefix}/restateTrain"
cache_config = CacheConfig(enable_caching=True, expire_after="30d")
xgb_image_uri = sagemaker.image_uris.retrieve(
framework="xgboost", # we are using the Sagemaker built in xgboost algorithm
region=region,
version="1.0-1",
py_version="py3",
instance_type=training_instance_type,
)
xgb_train = Estimator(
image_uri=xgb_image_uri,
instance_type=training_instance_type,
instance_count=1,
output_path=model_path,
base_job_name=f"{base_job_prefix}/restate-xgb-train",
sagemaker_session=sagemaker_session,
role=role,
)
xgb_train.set_hyperparameters(
# #objective="binary:logistic",
# objective="reg:linear",
num_round=50,
# max_depth=5,
# eta=0.2,
# gamma=4,
# min_child_weight=6,
# subsample=0.7,
# silent=0,
)
xgb_train.set_hyperparameters(grow_policy="lossguide")
xgb_objective_metric_name = "validation:mse"
xgb_hyperparameter_ranges = {
"max_depth": IntegerParameter(2, 10, scaling_type="Linear"),
}
xgb_tuner_log = HyperparameterTuner(
xgb_train,
xgb_objective_metric_name,
xgb_hyperparameter_ranges,
max_jobs=3,
max_parallel_jobs=3,
strategy="Random",
objective_type="Minimize",
)
xgb_step_tuning = TuningStep(
name="XGBHPTune",
tuner=xgb_tuner_log,
inputs={
"train": TrainingInput(
s3_data=step_process.properties.ProcessingOutputConfig.Outputs[
"train"
].S3Output.S3Uri,
content_type="text/csv",
),
"validation": TrainingInput(
s3_data=step_process.properties.ProcessingOutputConfig.Outputs[
"validation"
].S3Output.S3Uri,
content_type="text/csv",
),
},
cache_config=cache_config,
)
# dtree_image_uri = '625467769535.dkr.ecr.ap-southeast-1.amazonaws.com/sagemaker-decision-tree:latest'
dtree_image_uri = sagemaker_session.sagemaker_client.describe_image_version(
ImageName="restate-dtree"
)["ContainerImage"]
dtree_train = Estimator(
image_uri=dtree_image_uri,
role=role,
instance_count=1,
instance_type=training_instance_type,
base_job_name=f"{base_job_prefix}/restate-dtree-train",
output_path=model_path,
sagemaker_session=sagemaker_session,
)
dtree_objective_metric_name = "validation:mse"
dtree_metric_definitions = [{"Name": "validation:mse", "Regex": "mse:(\S+)"}]
dtree_hyperparameter_ranges = {
"max_depth": IntegerParameter(10, 50, scaling_type="Linear"),
"max_leaf_nodes": IntegerParameter(2, 12, scaling_type="Linear"),
}
dtree_tuner_log = HyperparameterTuner(
dtree_train,
dtree_objective_metric_name,
dtree_hyperparameter_ranges,
dtree_metric_definitions,
max_jobs=3,
max_parallel_jobs=3,
strategy="Random",
objective_type="Minimize",
)
dtree_step_tuning = TuningStep(
name="DTreeHPTune",
tuner=dtree_tuner_log,
inputs={
"training": TrainingInput(
s3_data=step_process.properties.ProcessingOutputConfig.Outputs[
"train"
].S3Output.S3Uri,
content_type="text/csv",
),
"validation": TrainingInput(
s3_data=step_process.properties.ProcessingOutputConfig.Outputs[
"validation"
].S3Output.S3Uri,
content_type="text/csv",
),
},
cache_config=cache_config,
)
dtree_script_eval = ScriptProcessor(
image_uri=dtree_image_uri,
command=["python3"],
instance_type=processing_instance_type,
instance_count=1,
base_job_name=f"{base_job_prefix}/script-dtree-eval",
sagemaker_session=sagemaker_session,
role=role,
)
dtree_evaluation_report = PropertyFile(
name="EvaluationReportDTree",
output_name="dtree_evaluation",
path="dtree_evaluation.json",
)
dtree_step_eval = ProcessingStep(
name="DTreeEval",
processor=dtree_script_eval,
inputs=[
ProcessingInput(
# source=dtree_step_train.properties.ModelArtifacts.S3ModelArtifacts,
source=dtree_step_tuning.get_top_model_s3_uri(top_k=0, s3_bucket=model_bucket_key),
destination="/opt/ml/processing/model",
),
ProcessingInput(
source=step_process.properties.ProcessingOutputConfig.Outputs[
"test"
].S3Output.S3Uri,
destination="/opt/ml/processing/test",
),
],
outputs=[
ProcessingOutput(
output_name="dtree_evaluation", source="/opt/ml/processing/evaluation"
),
],
code=os.path.join(BASE_DIR, "dtree_evaluate.py"),
property_files=[dtree_evaluation_report],
)
xgb_script_eval = ScriptProcessor(
image_uri=xgb_image_uri,
command=["python3"],
instance_type=processing_instance_type,
instance_count=1,
base_job_name=f"{base_job_prefix}/script-xgb-eval",
sagemaker_session=sagemaker_session,
role=role,
)
xgb_evaluation_report = PropertyFile(
name="EvaluationReportXGBoost",
output_name="xgb_evaluation",
path="xgb_evaluation.json",
)
xgb_step_eval = ProcessingStep(
name="XGBEval",
processor=xgb_script_eval,
inputs=[
ProcessingInput(
source=xgb_step_tuning.get_top_model_s3_uri(top_k=0, s3_bucket=model_bucket_key),
destination="/opt/ml/processing/model",
),
ProcessingInput(
source=step_process.properties.ProcessingOutputConfig.Outputs[
"test"
].S3Output.S3Uri,
destination="/opt/ml/processing/test",
),
],
outputs=[
ProcessingOutput(output_name="xgb_evaluation", source="/opt/ml/processing/evaluation"),
],
code=os.path.join(BASE_DIR, "xgb_evaluate.py"),
property_files=[xgb_evaluation_report],
)
xgb_model_metrics = ModelMetrics(
model_statistics=MetricsSource(
s3_uri="{}/xgb_evaluation.json".format(
xgb_step_eval.arguments["ProcessingOutputConfig"]["Outputs"][0]["S3Output"]["S3Uri"]
),
content_type="application/json",
)
)
dtree_model_metrics = ModelMetrics(
model_statistics=MetricsSource(
s3_uri="{}/dtree_evaluation.json".format(
dtree_step_eval.arguments["ProcessingOutputConfig"]["Outputs"][0]["S3Output"][
"S3Uri"
]
),
content_type="application/json",
)
)
xgb_eval_metrics = JsonGet(
step=xgb_step_eval,
property_file=xgb_evaluation_report,
json_path="regression_metrics.r2s.value", # This should follow the structure of your report_dict defined in the evaluate.py file.
)
dtree_eval_metrics = JsonGet(
step=dtree_step_eval,
property_file=dtree_evaluation_report,
json_path="regression_metrics.r2s.value", # This should follow the structure of your report_dict defined in the evaluate.py file.
)
# Register model step that will be conditionally executed
dtree_step_register = RegisterModel(
name="DTreeReg",
estimator=dtree_train,
model_data=dtree_step_tuning.get_top_model_s3_uri(top_k=0, s3_bucket=model_bucket_key),
content_types=["text/csv"],
response_types=["text/csv"],
inference_instances=["ml.t2.medium", "ml.m5.large"],
transform_instances=["ml.m5.large"],
model_package_group_name=model_package_group_name,
approval_status=model_approval_status,
model_metrics=dtree_model_metrics,
)
# Register model step that will be conditionally executed
xgb_step_register = RegisterModel(
name="XGBReg",
estimator=xgb_train,
model_data=xgb_step_tuning.get_top_model_s3_uri(top_k=0, s3_bucket=model_bucket_key),
content_types=["text/csv"],
response_types=["text/csv"],
inference_instances=["ml.t2.medium", "ml.m5.large"],
transform_instances=["ml.m5.large"],
model_package_group_name=model_package_group_name,
approval_status=model_approval_status,
model_metrics=xgb_model_metrics,
)
# Condition step for evaluating model quality and branching execution
cond_lte = ConditionGreaterThanOrEqualTo( # You can change the condition here
left=JsonGet(
step=dtree_step_eval,
property_file=dtree_evaluation_report,
json_path="regression_metrics.r2s.value", # This should follow the structure of your report_dict defined in the evaluate.py file.
),
right=JsonGet(
step=xgb_step_eval,
property_file=xgb_evaluation_report,
json_path="regression_metrics.r2s.value", # This should follow the structure of your report_dict defined in the evaluate.py file.
), # You can change the threshold here
)
step_cond = ConditionStep(
name="AccuracyCond",
conditions=[cond_lte],
if_steps=[dtree_step_register],
else_steps=[xgb_step_register],
)
create_date = time.strftime("%Y-%m-%d-%H-%M-%S")
# Pipeline instance
pipeline = Pipeline(
name=pipeline_name,
parameters=[
processing_instance_type,
processing_instance_count,
training_instance_type,
model_approval_status,
input_data
],
pipeline_experiment_config=PipelineExperimentConfig(
pipeline_name + "-" + create_date, "restate-{}".format(create_date)
),
steps=[
data_wrangler_step,
step_process,
dtree_step_tuning,
xgb_step_tuning,
dtree_step_eval,
xgb_step_eval,
step_cond,
],
sagemaker_session=sagemaker_session,
)
return pipeline
ROLE = "arn:aws:iam::012345678901:role/SageMakerRole"
INSTANCE_COUNT = 1
INSTANCE_TYPE = "ml.m5.10xlarge"
VOLUME_SIZE_IN_GB = 2
VOLUME_KMS_KEY = "volume-kms-key"
OUTPUT_KMS_KEY = "output-kms-key"
MAX_RUNTIME_IN_SECONDS = 3
BASE_JOB_NAME = "base-job-name"
ENV_KEY_1 = "env_key_1"
ENV_VALUE_1 = "env_key_1"
ENVIRONMENT = {ENV_KEY_1: ENV_VALUE_1}
TAG_KEY_1 = "tag_key_1"
TAG_VALUE_1 = "tag_value_1"
TAGS = [{"Key": TAG_KEY_1, "Value": TAG_VALUE_1}]
NETWORK_CONFIG = NetworkConfig(enable_network_isolation=True)
ENABLE_CLOUDWATCH_METRICS = True
BASELINE_DATASET_PATH = "/my/local/path/baseline.csv"
PREPROCESSOR_PATH = "/my/local/path/preprocessor.py"
POSTPROCESSOR_PATH = "/my/local/path/postprocessor.py"
OUTPUT_S3_URI = "s3://output-s3-uri/"
CUSTOM_IMAGE_URI = "012345678901.dkr.ecr.us-west-2.amazonaws.com/my-custom-image-uri"
DEFAULT_IMAGE_URI = "159807026194.dkr.ecr.us-west-2.amazonaws.com/sagemaker-model-monitor-analyzer"
INTER_CONTAINER_ENCRYPTION_EXCEPTION_MSG = (
"EnableInterContainerTrafficEncryption is not supported in Model Monitor. Please ensure that "
)
"encrypt_inter_container_traffic=None when creating your NetworkConfig object."
def get_pipeline(
region,
security_group_ids,
subnets,
processing_role=None,
training_role=None,
data_bucket=None,
model_bucket=None,
model_package_group_name="AbalonePackageGroup",
pipeline_name="AbalonePipeline",
base_job_prefix="Abalone",
):
"""Gets a SageMaker ML Pipeline instance working with on abalone data.
Args:
region: AWS region to create and run the pipeline.
processing_role: IAM role to create and run processing steps
training_role: IAM role to create and run training steps
data_bucket: the bucket to use for storing the artifacts
Returns:
an instance of a pipeline
"""
sagemaker_session = get_session(region, data_bucket)
if processing_role is None:
processing_role = sagemaker.session.get_execution_role(sagemaker_session)
if training_role is None:
training_role = sagemaker.session.get_execution_role(sagemaker_session)
if model_bucket is None:
model_bucket = sagemaker_session.default_bucket()
print(f"Creating the pipeline '{pipeline_name}':")
print(f"Parameters:{region}\n{security_group_ids}\n{subnets}\n{processing_role}\n\
{training_role}\n{data_bucket}\n{model_bucket}\n{model_package_group_name}\n\
{pipeline_name}\n{base_job_prefix}")
# parameters for pipeline execution
processing_instance_count = ParameterInteger(name="ProcessingInstanceCount", default_value=1)
processing_instance_type = ParameterString(
name="ProcessingInstanceType", default_value="ml.m5.xlarge"
)
training_instance_type = ParameterString(
name="TrainingInstanceType", default_value="ml.m5.xlarge"
)
model_approval_status = ParameterString(
name="ModelApprovalStatus", default_value="PendingManualApproval"
)
input_data = ParameterString(
name="InputDataUrl",
default_value=f"s3://{sagemaker_session.default_bucket()}/datasets/abalone-dataset.csv",
)
# configure network for encryption, network isolation and VPC configuration
# Since the preprocessor job takes the data from S3, enable_network_isolation must be set to False
# see https://github.com/aws/amazon-sagemaker-examples/issues/1689
network_config = NetworkConfig(
enable_network_isolation=False,
security_group_ids=security_group_ids.split(","),
subnets=subnets.split(","),
encrypt_inter_container_traffic=True)
# processing step for feature engineering
sklearn_processor = SKLearnProcessor(
framework_version="0.23-1",
instance_type=processing_instance_type,
instance_count=processing_instance_count,
base_job_name=f"{base_job_prefix}/sklearn-abalone-preprocess",
sagemaker_session=sagemaker_session,
role=processing_role,
network_config=network_config
)
step_process = ProcessingStep(
name="PreprocessAbaloneData",
processor=sklearn_processor,
outputs=[
ProcessingOutput(output_name="train", source="/opt/ml/processing/train"),
ProcessingOutput(output_name="validation", source="/opt/ml/processing/validation"),
ProcessingOutput(output_name="test", source="/opt/ml/processing/test"),
],
code=os.path.join(BASE_DIR, "preprocess.py"),
job_arguments=["--input-data", input_data],
)
# training step for generating model artifacts
model_path = f"s3://{model_bucket}/{base_job_prefix}/AbaloneTrain"
image_uri = sagemaker.image_uris.retrieve(
framework="xgboost",
region=region,
version="1.0-1",
py_version="py3",
instance_type=training_instance_type,
)
xgb_train = Estimator(
image_uri=image_uri,
instance_type=training_instance_type,
instance_count=1,
output_path=model_path,
base_job_name=f"{base_job_prefix}/abalone-train",
sagemaker_session=sagemaker_session,
role=training_role,
subnets=network_config.subnets,
security_group_ids=network_config.security_group_ids,
encrypt_inter_container_traffic=True,
enable_network_isolation=False
)
xgb_train.set_hyperparameters(
objective="reg:linear",
num_round=50,
max_depth=5,
eta=0.2,
gamma=4,
min_child_weight=6,
subsample=0.7,
silent=0,
)
step_train = TrainingStep(
name="TrainAbaloneModel",
estimator=xgb_train,
inputs={
"train": TrainingInput(
s3_data=step_process.properties.ProcessingOutputConfig.Outputs[
"train"
].S3Output.S3Uri,
content_type="text/csv",
),
"validation": TrainingInput(
s3_data=step_process.properties.ProcessingOutputConfig.Outputs[
"validation"
].S3Output.S3Uri,
content_type="text/csv",
),
},
)
# processing step for evaluation
script_eval = ScriptProcessor(
image_uri=image_uri,
command=["python3"],
instance_type=processing_instance_type,
instance_count=1,
base_job_name=f"{base_job_prefix}/script-abalone-eval",
sagemaker_session=sagemaker_session,
role=processing_role,
network_config=network_config
)
evaluation_report = PropertyFile(
name="AbaloneEvaluationReport",
output_name="evaluation",
path="evaluation.json",
)
step_eval = ProcessingStep(
name="EvaluateAbaloneModel",
processor=script_eval,
inputs=[
ProcessingInput(
source=step_train.properties.ModelArtifacts.S3ModelArtifacts,
destination="/opt/ml/processing/model",
),
ProcessingInput(
source=step_process.properties.ProcessingOutputConfig.Outputs[
"test"
].S3Output.S3Uri,
destination="/opt/ml/processing/test",
),
],
outputs=[
ProcessingOutput(output_name="evaluation", source="/opt/ml/processing/evaluation"),
],
code=os.path.join(BASE_DIR, "evaluate.py"),
property_files=[evaluation_report],
)
# register model step that will be conditionally executed
model_metrics = ModelMetrics(
model_statistics=MetricsSource(
s3_uri="{}/evaluation.json".format(
step_eval.arguments["ProcessingOutputConfig"]["Outputs"][0]["S3Output"]["S3Uri"]
),
content_type="application/json"
)
)
"""
There is a bug in RegisterModel implementation
The RegisterModel step is implemented in the SDK as two steps, a _RepackModelStep and a _RegisterModelStep.
The _RepackModelStep runs a SKLearn training step in order to repack the model.tar.gz to include any custom inference code in the archive.
The _RegisterModelStep then registers the repacked model.
The problem is that the _RepackModelStep does not propagate VPC configuration from the Estimator object:
https://github.com/aws/sagemaker-python-sdk/blob/cdb633b3ab02398c3b77f5ecd2c03cdf41049c78/src/sagemaker/workflow/_utils.py#L88
This cause the AccessDenied exception because repacker cannot access S3 bucket (all access which is not via VPC endpoint is bloked by the bucket policy)
The issue is opened against SageMaker python SDK: https://github.com/aws/sagemaker-python-sdk/issues/2302
"""
vpc_config = {
"Subnets":network_config.subnets,
"SecurityGroupIds":network_config.security_group_ids
}
step_register = RegisterModel(
name="RegisterAbaloneModel",
estimator=xgb_train,
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"],
model_package_group_name=model_package_group_name,
approval_status=model_approval_status,
model_metrics=model_metrics,
vpc_config_override=vpc_config
)
# condition step for evaluating model quality and branching execution
cond_lte = ConditionLessThanOrEqualTo(
left=JsonGet(
step=step_eval,
property_file=evaluation_report,
json_path="regression_metrics.mse.value"
),
right=6.0,
)
step_cond = ConditionStep(
name="CheckMSEAbaloneEvaluation",
conditions=[cond_lte],
if_steps=[step_register],
else_steps=[],
)
# pipeline instance
pipeline = Pipeline(
name=pipeline_name,
parameters=[
processing_instance_type,
processing_instance_count,
training_instance_type,
model_approval_status,
input_data,
],
steps=[step_process, step_train, step_eval, step_cond],
sagemaker_session=sagemaker_session,
)
return pipeline