def test_data_wrangler_processor_with_required_parameters(sagemaker_session):
processor = DataWranglerProcessor(
role=ROLE,
data_wrangler_flow_source=DATA_WRANGLER_RECIPE_SOURCE,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
)
processor.run()
expected_args = _get_expected_args(processor._current_job_name)
sagemaker_session.process.assert_called_with(**expected_args)
def test_data_wrangler_processor_with_mock_input(sagemaker_session):
processor = DataWranglerProcessor(
role=ROLE,
data_wrangler_flow_source=DATA_WRANGLER_RECIPE_SOURCE,
instance_count=1,
instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
)
mock_input = ProcessingInput(
source=MOCK_S3_URI,
destination="/opt/ml/processing/mock_input",
input_name="mock_input",
s3_data_type="S3Prefix",
s3_input_mode="File",
s3_data_distribution_type="FullyReplicated",
)
processor.run(inputs=[mock_input])
expected_args = _get_expected_args(processor._current_job_name,
add_mock_input=True)
sagemaker_session.process.assert_called_with(**expected_args)
def test_one_step_data_wrangler_processing_pipeline(sagemaker_session, role,
pipeline_name):
instance_count = ParameterInteger(name="InstanceCount", default_value=1)
instance_type = ParameterString(name="InstanceType",
default_value="ml.m5.4xlarge")
recipe_file_path = os.path.join(DATA_DIR, "workflow", "dummy_recipe.flow")
input_file_path = os.path.join(DATA_DIR, "workflow", "dummy_data.csv")
output_name = "3f74973c-fd1e-4845-89f8-0dd400031be9.default"
output_content_type = "CSV"
output_config = {output_name: {"content_type": output_content_type}}
job_argument = [f"--output-config '{json.dumps(output_config)}'"]
inputs = [
ProcessingInput(
input_name="dummy_data.csv",
source=input_file_path,
destination="/opt/ml/processing/dummy_data.csv",
)
]
output_s3_uri = f"s3://{sagemaker_session.default_bucket()}/output"
outputs = [
ProcessingOutput(
output_name=output_name,
source="/opt/ml/processing/output",
destination=output_s3_uri,
s3_upload_mode="EndOfJob",
)
]
data_wrangler_processor = DataWranglerProcessor(
role=role,
data_wrangler_flow_source=recipe_file_path,
instance_count=instance_count,
instance_type=instance_type,
sagemaker_session=sagemaker_session,
max_runtime_in_seconds=86400,
)
data_wrangler_step = ProcessingStep(
name="data-wrangler-step",
processor=data_wrangler_processor,
inputs=inputs,
outputs=outputs,
job_arguments=job_argument,
)
pipeline = Pipeline(
name=pipeline_name,
parameters=[instance_count, instance_type],
steps=[data_wrangler_step],
sagemaker_session=sagemaker_session,
)
definition = json.loads(pipeline.definition())
expected_image_uri = image_uris.retrieve(
"data-wrangler", region=sagemaker_session.boto_region_name)
assert len(definition["Steps"]) == 1
assert definition["Steps"][0]["Arguments"]["AppSpecification"][
"ImageUri"] is not None
assert definition["Steps"][0]["Arguments"]["AppSpecification"][
"ImageUri"] == expected_image_uri
assert definition["Steps"][0]["Arguments"]["ProcessingInputs"] is not None
processing_inputs = definition["Steps"][0]["Arguments"]["ProcessingInputs"]
assert len(processing_inputs) == 2
for processing_input in processing_inputs:
if processing_input["InputName"] == "flow":
assert processing_input["S3Input"]["S3Uri"].endswith(".flow")
assert processing_input["S3Input"][
"LocalPath"] == "/opt/ml/processing/flow"
elif processing_input["InputName"] == "dummy_data.csv":
assert processing_input["S3Input"]["S3Uri"].endswith(".csv")
assert processing_input["S3Input"][
"LocalPath"] == "/opt/ml/processing/dummy_data.csv"
else:
raise AssertionError("Unknown input name")
assert definition["Steps"][0]["Arguments"][
"ProcessingOutputConfig"] is not None
processing_outputs = definition["Steps"][0]["Arguments"][
"ProcessingOutputConfig"]["Outputs"]
assert len(processing_outputs) == 1
assert processing_outputs[0]["OutputName"] == output_name
assert processing_outputs[0]["S3Output"] is not None
assert processing_outputs[0]["S3Output"][
"LocalPath"] == "/opt/ml/processing/output"
assert processing_outputs[0]["S3Output"]["S3Uri"] == output_s3_uri
try:
response = pipeline.create(role)
create_arn = response["PipelineArn"]
execution = pipeline.start()
response = execution.describe()
assert response["PipelineArn"] == create_arn
try:
execution.wait(delay=60, max_attempts=10)
except WaiterError:
pass
execution_steps = execution.list_steps()
assert len(execution_steps) == 1
assert execution_steps[0]["StepName"] == "data-wrangler-step"
finally:
try:
pipeline.delete()
except Exception:
pass
),
(
MXNetProcessor(
framework_version="1.4.1",
py_version="py3",
role=sagemaker.get_execution_role(),
instance_count=1,
instance_type=INSTANCE_TYPE,
base_job_name="test-mxnet",
),
{"code": DUMMY_S3_SCRIPT_PATH},
),
(
DataWranglerProcessor(
role=sagemaker.get_execution_role(),
data_wrangler_flow_source="s3://my-bucket/dw.flow",
instance_count=1,
instance_type=INSTANCE_TYPE,
),
{},
),
(
SparkJarProcessor(
role=sagemaker.get_execution_role(),
framework_version="2.4",
instance_count=1,
instance_type=INSTANCE_TYPE,
),
{
"submit_app": "s3://my-jar",
"submit_class": "com.amazonaws.sagemaker.spark.test.HelloJavaSparkApp",
"arguments": ["--input", "input-data-uri", "--output", "output-data-uri"],
def test_one_step_ingestion_pipeline(sagemaker_session, feature_store_session,
feature_definitions, role, pipeline_name):
instance_count = ParameterInteger(name="InstanceCount", default_value=1)
instance_type = ParameterString(name="InstanceType",
default_value="ml.m5.4xlarge")
input_name = "features.csv"
input_file_path = os.path.join(DATA_DIR, "workflow", "features.csv")
input_data_uri = os.path.join("s3://", sagemaker_session.default_bucket(),
"py-sdk-ingestion-test-input/features.csv")
with open(input_file_path, "r") as data:
body = data.read()
S3Uploader.upload_string_as_file_body(
body=body,
desired_s3_uri=input_data_uri,
sagemaker_session=sagemaker_session)
inputs = [
ProcessingInput(
input_name=input_name,
source=input_data_uri,
destination="/opt/ml/processing/features.csv",
)
]
feature_group_name = f"py-sdk-integ-fg-{int(time.time() * 10**7)}"
feature_group = FeatureGroup(
name=feature_group_name,
feature_definitions=feature_definitions,
sagemaker_session=feature_store_session,
)
ingestion_only_flow, output_name = generate_data_ingestion_flow_from_s3_input(
input_name,
input_data_uri,
s3_content_type="csv",
s3_has_header=True,
)
outputs = [
ProcessingOutput(
output_name=output_name,
app_managed=True,
feature_store_output=FeatureStoreOutput(
feature_group_name=feature_group_name),
)
]
temp_flow_path = "./ingestion.flow"
with cleanup_feature_group(feature_group):
json.dump(ingestion_only_flow, open(temp_flow_path, "w"))
data_wrangler_processor = DataWranglerProcessor(
role=role,
data_wrangler_flow_source=temp_flow_path,
instance_count=instance_count,
instance_type=instance_type,
sagemaker_session=sagemaker_session,
max_runtime_in_seconds=86400,
)
data_wrangler_step = ProcessingStep(name="ingestion-step",
processor=data_wrangler_processor,
inputs=inputs,
outputs=outputs)
pipeline = Pipeline(
name=pipeline_name,
parameters=[instance_count, instance_type],
steps=[data_wrangler_step],
sagemaker_session=sagemaker_session,
)
try:
response = pipeline.create(role)
create_arn = response["PipelineArn"]
offline_store_s3_uri = os.path.join(
"s3://", sagemaker_session.default_bucket(),
feature_group_name)
feature_group.create(
s3_uri=offline_store_s3_uri,
record_identifier_name="f11",
event_time_feature_name="f10",
role_arn=role,
enable_online_store=False,
)
_wait_for_feature_group_create(feature_group)
execution = pipeline.start()
response = execution.describe()
assert response["PipelineArn"] == create_arn
try:
execution.wait(delay=60, max_attempts=10)
except WaiterError:
pass
execution_steps = execution.list_steps()
assert len(execution_steps) == 1
assert execution_steps[0]["StepName"] == "ingestion-step"
assert execution_steps[0]["StepStatus"] == "Succeeded"
athena_query = feature_group.athena_query()
with timeout(minutes=10):
athena_query.run(
query_string=f'SELECT * FROM "{athena_query.table_name}"',
output_location=f"{offline_store_s3_uri}/query_results",
)
athena_query.wait()
assert "SUCCEEDED" == athena_query.get_query_execution().get(
"QueryExecution").get("Status").get("State")
df = athena_query.as_dataframe()
assert pd.read_csv(input_file_path).shape[0] == df.shape[0]
finally:
try:
pipeline.delete()
except Exception as e:
print(f"Delete pipeline failed with error: {e}")
os.remove(temp_flow_path)